From c49ec7e20c963210762ee68f1cb3dd1e6b11eb60 Mon Sep 17 00:00:00 2001 From: Janosch <99879757+jkppr@users.noreply.github.com> Date: Wed, 27 Nov 2024 18:59:18 +0100 Subject: [PATCH] Fix: Resolve pytype --strict-none-binding issue in the api client (#3214) * Fixing pytype strict-none-binding issue. --------- Co-authored-by: Thomas Chopitea --- .../python/timesketch_api_client/search.py | 30 +++++++++++++------ .../python/timesketch_api_client/story.py | 3 ++ 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/api_client/python/timesketch_api_client/search.py b/api_client/python/timesketch_api_client/search.py index 2663b60d37..027ffbe870 100644 --- a/api_client/python/timesketch_api_client/search.py +++ b/api_client/python/timesketch_api_client/search.py @@ -493,12 +493,19 @@ def _execute_query(self, file_name="", count=False): """Execute a search request and store the results. Args: - file_name (str): optional file path to a filename that + file_name (str): Optional file path to a filename that all the results will be saved to. If not provided the results will be stored in the search object. - count (bool): optional boolean that determines whether + count (bool): Optional boolean that determines whether we want to execute the query or only count the - number of events that the query would produce. + number of events that the query would produce. If + set to True, the results will be stored in the + search object, and the number of events will be + returned. + + Returns: + A dict with the search results or the total number of events + (if count=True) or None if saved to file. """ query_filter = self.query_filter if not isinstance(query_filter, dict): @@ -531,14 +538,14 @@ def _execute_query(self, file_name="", count=False): if file_name: with open(file_name, "wb") as fw: fw.write(response.content) - return + return None response_json = error.get_response_json(response, logger) if count: meta = response_json.get("meta", {}) self._total_elastic_size = meta.get("total_count", 0) - return + return meta.get("total_count", 0) scroll_id = response_json.get("meta", {}).get("scroll_id", "") form_data["scroll_id"] = scroll_id @@ -579,6 +586,7 @@ def _execute_query(self, file_name="", count=False): ) self._raw_response = response_json + return response_json def add_chip(self, chip): """Add a chip to the ...""" @@ -647,7 +655,7 @@ def expected_size(self): if self._total_elastic_size: return self._total_elastic_size - self._execute_query(count=True) + _ = self._execute_query(count=True) return self._total_elastic_size def from_manual( # pylint: disable=arguments-differ @@ -1074,8 +1082,10 @@ def scrolling_enable(self): def to_dict(self): """Returns a dict with the respone of the query.""" - if not self._raw_response: + if self._raw_response is None: self._execute_query() + if self._raw_response is None: + raise ValueError("No results to return.") return self._raw_response @@ -1098,8 +1108,10 @@ def to_file(self, file_name): def to_pandas(self): """Returns a pandas DataFrame with the response of the query.""" - if not self._raw_response: - self._execute_query() + if self._raw_response is None: + self._raw_response = self._execute_query() + if self._raw_response is None: + raise ValueError("No results to return.") return_list = [] timelines = {t.id: t.name for t in self._sketch.list_timelines()} diff --git a/api_client/python/timesketch_api_client/story.py b/api_client/python/timesketch_api_client/story.py index 92dc0e55cf..6bf49faf3f 100644 --- a/api_client/python/timesketch_api_client/story.py +++ b/api_client/python/timesketch_api_client/story.py @@ -729,6 +729,9 @@ def to_string(self): string_list.append(block.text) elif block.TYPE == "view": search_obj = block.view + if search_obj is None: + logging.warning("Block has no view. Skipping") + continue data_frame = search_obj.to_pandas() string_list.append(data_frame.to_string(index=False)) elif block.TYPE == "aggregation":