chore: add function return type annotations

Author: mpetazzoni

wherobots/db/connection.py

@@ -74,19 +74,19 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def close(self):
+    def close(self) -> None:
         self.__ws.close()
 
-    def commit(self):
+    def commit(self) -> None:
         raise NotSupportedError
 
-    def rollback(self):
+    def rollback(self) -> None:
         raise NotSupportedError
 
     def cursor(self) -> Cursor:
         return Cursor(self.__execute_sql, self.__cancel_query)
 
-    def __main_loop(self):
+    def __main_loop(self) -> None:
         """Main background loop listening for messages from the SQL session."""
         logging.info("Starting background connection handling loop...")
         while self.__ws.protocol.state < websockets.protocol.State.CLOSING:
@@ -101,7 +101,7 @@ def __main_loop(self):
             except Exception as e:
                 logging.exception("Error handling message from SQL session", exc_info=e)
 
-    def __listen(self):
+    def __listen(self) -> None:
         """Waits for the next message from the SQL session and processes it.
 
         The code in this method is purposefully defensive to avoid unexpected situations killing the thread.

wherobots/db/constants.py

@@ -44,7 +44,7 @@ class ExecutionState(LowercaseStrEnum):
     COMPLETED = auto()
     "The driver has completed processing the query results."
 
-    def is_terminal_state(self):
+    def is_terminal_state(self) -> bool:
         return self in (
             ExecutionState.COMPLETED,
             ExecutionState.CANCELLED,
@@ -97,7 +97,7 @@ class AppStatus(StrEnum):
     DESTROY_FAILED = auto()
     DESTROYED = auto()
 
-    def is_starting(self):
+    def is_starting(self) -> bool:
         return self in (
             AppStatus.PENDING,
             AppStatus.PREPARING,
@@ -107,7 +107,7 @@ def is_starting(self):
             AppStatus.INITIALIZING,
         )
 
-    def is_terminal_state(self):
+    def is_terminal_state(self) -> bool:
         return self in (
             AppStatus.PREPARE_FAILED,
             AppStatus.DEPLOY_FAILED,

wherobots/db/cursor.py

@@ -16,7 +16,7 @@
 
 class Cursor:
 
-    def __init__(self, exec_fn, cancel_fn):
+    def __init__(self, exec_fn, cancel_fn) -> None:
         self.__exec_fn = exec_fn
         self.__cancel_fn = cancel_fn
 
@@ -72,7 +72,7 @@ def __get_results(self) -> Optional[List[Tuple[Any, ...]]]:
 
         return self.__results
 
-    def execute(self, operation: str, parameters: dict[str, Any] = None):
+    def execute(self, operation: str, parameters: dict[str, Any] = None) -> None:
         if self.__current_execution_id:
             self.__cancel_fn(self.__current_execution_id)
 
@@ -84,38 +84,40 @@ def execute(self, operation: str, parameters: dict[str, Any] = None):
         sql = operation.format(**(parameters or {}))
         self.__current_execution_id = self.__exec_fn(sql, self.__on_execution_result)
 
-    def executemany(self, operation: str, seq_of_parameters: list[dict[str, Any]]):
+    def executemany(
+        self, operation: str, seq_of_parameters: list[dict[str, Any]]
+    ) -> None:
         raise NotImplementedError
 
-    def fetchone(self):
+    def fetchone(self) -> Any:
         results = self.__get_results()[self.__current_row :]
         if len(results) == 0:
             return None
         self.__current_row += 1
         return results[0]
 
-    def fetchmany(self, size: int = None):
+    def fetchmany(self, size: int = None) -> list[Any]:
         size = size or self.arraysize
         results = self.__get_results()[self.__current_row : self.__current_row + size]
         self.__current_row += size
         return results
 
-    def fetchall(self):
+    def fetchall(self) -> list[Any]:
         return self.__get_results()[self.__current_row :]
 
-    def close(self):
+    def close(self) -> None:
         """Close the cursor."""
         if self.__results is None and self.__current_execution_id:
             self.__cancel_fn(self.__current_execution_id)
 
-    def __iter__(self):
+    def __iter__(self) -> Cursor:
         return self
 
-    def __next__(self):
+    def __next__(self) -> None:
         raise StopIteration
 
-    def __enter__(self):
+    def __enter__(self) -> Cursor:
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.close()