fix(api-client): code hygiene

api_client/python/timesketch_api_client/aggregation_test.py

@@ -14,7 +14,7 @@
 """Tests for the Timesketch API aggregation object."""
 
 import unittest
-import mock
+from unittest import mock
 
 import altair as alt
 

api_client/python/timesketch_api_client/analyzer.py

@@ -13,23 +13,36 @@
 # limitations under the License.
 """Timesketch API analyzer result object."""
 
-from __future__ import unicode_literals
+from __future__ import annotations
 
 import datetime
 import json
 import logging
+from typing import Any, Dict, Generator, List, TYPE_CHECKING
 
 from . import error
 from . import resource
 
+if TYPE_CHECKING:
+    from .client import TimesketchApi
+
 logger = logging.getLogger("timesketch_api.analyzer")
 
 
 class AnalyzerResult(resource.BaseResource):
     """Class to store and retrieve session information for an analyzer."""
 
-    def __init__(self, timeline_id, session_id, sketch_id, api):
-        """Initialize the class."""
+    def __init__(
+        self, timeline_id: int, session_id: int, sketch_id: int, api: TimesketchApi
+    ) -> None:
+        """Initialize the class.
+
+        Args:
+            timeline_id (int): The ID of the timeline.
+            session_id (int): The ID of the analyzer session.
+            sketch_id (int): The ID of the sketch.
+            api (TimesketchApi): An instance of TimesketchApi.
+        """
         self._session_id = session_id
         self._sketch_id = sketch_id
         self._timeline_id = timeline_id
@@ -38,21 +51,21 @@ def __init__(self, timeline_id, session_id, sketch_id, api):
         )
         super().__init__(api, resource_uri)
 
-    def _get_status_data(self):
+    def _get_status_data(self) -> Generator[Dict[str, str], None, None]:
         """Yields a dict for each analyzer status."""
         data = self._fetch_data()
         for entry in data.get("analyzers", []):
             yield {
-                "log": entry.get("log", "No recorded logs."),
-                "name": entry.get("name", "No Name"),
-                "results": entry.get("results", ""),
-                "status": entry.get("status", "Unknown"),
-                "date": entry.get(
-                    "status_date", datetime.datetime.utcnow().isoformat()
+                "log": str(entry.get("log", "No recorded logs.")),
+                "name": str(entry.get("name", "No Name")),
+                "results": str(entry.get("results", "")),
+                "status": str(entry.get("status", "Unknown")),
+                "date": str(
+                    entry.get("status_date", datetime.datetime.utcnow().isoformat())
                 ),
             }
 
-    def _fetch_data(self):
+    def _fetch_data(self) -> Dict[str, Any]:
         """Returns a dict with the analyzer results."""
         response = self.api.session.get(self.resource_uri)
         if not error.check_return_status(response, logger):
@@ -64,7 +77,7 @@ def _fetch_data(self):
         if not objects:
             return {}
 
-        result_dict = {}
+        result_dict: Dict[str, Any] = {}
         for result in objects[0]:
             result_id = result.get("analysissession_id")
             if result_id != self._session_id:
@@ -100,14 +113,14 @@ def _fetch_data(self):
         return result_dict
 
     @property
-    def id(self):
+    def id(self) -> int:
         """Returns the session ID."""
         return self._session_id
 
     @property
-    def log(self):
+    def log(self) -> str:
         """Returns back logs from the analyzer session, if there are any."""
-        return_strings = []
+        return_strings: List[str] = []
         for entry in self._get_status_data():
             return_strings.append(
                 "[{0:s}] = {1:s}".format(
@@ -118,9 +131,9 @@ def log(self):
         return "\n".join(return_strings)
 
     @property
-    def results(self):
+    def results(self) -> str:
         """Returns the results from the analyzer session."""
-        return_strings = []
+        return_strings: List[str] = []
         for entry in self._get_status_data():
             results = entry.get("results")
             if not results:
@@ -131,9 +144,9 @@ def results(self):
         return "\n".join(return_strings)
 
     @property
-    def results_dict(self):
+    def results_dict(self) -> Dict[str, List[str]]:
         """Returns the results from the analyzer session as a dict."""
-        result_dict = {}
+        result_dict: Dict[str, List[str]] = {}
         for entry in self._get_status_data():
             results = entry.get("results")
             if not results:
@@ -144,9 +157,9 @@ def results_dict(self):
         return result_dict
 
     @property
-    def status(self):
+    def status(self) -> str:
         """Returns the current status of the analyzer run."""
-        return_strings = []
+        return_strings: List[str] = []
         for entry in self._get_status_data():
             return_strings.append(
                 "[{0:s}] = {1:s}".format(
@@ -157,9 +170,9 @@ def status(self):
         return "\n".join(return_strings)
 
     @property
-    def status_dict(self):
+    def status_dict(self) -> Dict[str, List[str]]:
         """Returns the current status of the analyzers run as a dict."""
-        return_dict = {}
+        return_dict: Dict[str, List[str]] = {}
         for entry in self._get_status_data():
             name = entry.get("name", "No Name")
 
@@ -168,9 +181,9 @@ def status_dict(self):
         return return_dict
 
     @property
-    def status_string(self):
+    def status_string(self) -> str:
         """Returns a longer version of a status string."""
-        return_strings = []
+        return_strings: List[str] = []
         for entry in self._get_status_data():
             return_strings.append(
                 "{0:s} - {1:s}: {2:s}".format(

api_client/python/timesketch_api_client/cli_input.py

@@ -13,26 +13,24 @@
 # limitations under the License.
 """CLI assistance for importer tools."""
 
-from typing import Any
-from typing import Callable
-from typing import Optional
-from typing import Text
+from __future__ import annotations
 
 import getpass
+from typing import Any, Callable, Optional
 
 
 def ask_question(
-    question: Text,
-    input_type: Callable[[Text], Any],
+    question: str,
+    input_type: Callable[[str], Any],
     default: Optional[Any] = None,
-    hide_input: Optional[bool] = False,
+    hide_input: bool = False,
 ) -> Any:
     """Presents the user with a prompt with a default return value and a type.
 
     Args:
         question (str): the text that the user will be prompted.
         input_type (type): the type of the input data.
-        default (object): default value for the question, optional.
+        default: default value for the question, optional.
         hide_input (bool): whether the input should be hidden, eg. when asking
             for a password.
 
@@ -56,9 +54,7 @@ def ask_question(
     return input_type(answer)
 
 
-def confirm_choice(
-    choice: Text, default: Optional[bool] = True, abort: Optional[bool] = True
-) -> bool:
+def confirm_choice(choice: str, default: bool = True, abort: bool = True) -> bool:
     """Returns a bool from a yes/no question presented to the end user.
 
     Args: