Add typing hints to session managers and handler (#1135)

* Add typing hints to session managers and handler

* Explicitly return None from start_session
* Update MappingKernelManager to RemoteMappingKernelManager
* Use __future__ to resolve circular import
* Remove unused imports
* Remove import statement for RemoteMappingKernelManager and add noqa F821

Author: kiersten-stokes

enterprise_gateway/services/sessions/handlers.py

@@ -17,7 +17,7 @@ class SessionRootHandler(
     JSON errors.
     """
 
-    async def get(self):
+    async def get(self) -> None:
         """Overrides the super class method to honor the kernel listing
         configuration setting.
 

enterprise_gateway/services/sessions/kernelsessionmanager.py

@@ -2,6 +2,8 @@
 # Distributed under the terms of the Modified BSD License.
 """Session manager that keeps all its metadata in memory."""
 
+from __future__ import annotations
+
 import getpass
 import json
 import os
@@ -49,7 +51,7 @@ class KernelSessionManager(LoggingConfigurable):
     )
 
     @default("enable_persistence")
-    def session_persistence_default(self):
+    def session_persistence_default(self) -> bool:
         return bool(
             os.getenv(
                 self.session_persistence_env, str(self.session_persistence_default_value)
@@ -66,16 +68,16 @@ def session_persistence_default(self):
     )
 
     @default("persistence_root")
-    def persistence_root_default(self):
+    def persistence_root_default(self) -> str:
         return os.getenv(self.persistence_root_env, "/")
 
-    def __init__(self, kernel_manager, **kwargs):
+    def __init__(self, kernel_manager: RemoteMappingKernelManager, **kwargs):  # noqa: F821
         super().__init__(**kwargs)
         self.kernel_manager = kernel_manager
         self._sessions = dict()
         self._sessionsByUser = dict()
 
-    def create_session(self, kernel_id, **kwargs):
+    def create_session(self, kernel_id: str, **kwargs) -> None:
         """
         Creates a session associated with this kernel.
 
@@ -106,7 +108,7 @@ def create_session(self, kernel_id, **kwargs):
         )
         self._save_session(kernel_id, kernel_session)
 
-    def refresh_session(self, kernel_id):
+    def refresh_session(self, kernel_id: str) -> None:
         """
         Refreshes the session from its persisted state. Called on kernel restarts.
         """
@@ -123,7 +125,7 @@ def refresh_session(self, kernel_id):
         )
         self._save_session(kernel_id, kernel_session)
 
-    def _save_session(self, kernel_id, kernel_session):
+    def _save_session(self, kernel_id: str, kernel_session: dict) -> None:
         # Write/commit the addition, update dictionary
         kernels_lock.acquire()
         try:
@@ -140,12 +142,13 @@ def _save_session(self, kernel_id, kernel_session):
         finally:
             kernels_lock.release()
 
-    def start_session(self, kernel_id):
+    def start_session(self, kernel_id: str) -> bool | None:
         kernel_session = self._sessions.get(kernel_id, None)
         if kernel_session is not None:
             return self._start_session(kernel_session)
+        return None
 
-    def start_sessions(self):
+    def start_sessions(self) -> None:
         """
         Attempt to start persisted sessions.
 
@@ -173,7 +176,7 @@ def start_sessions(self):
 
             self._delete_sessions(sessions_to_remove)
 
-    def _start_session(self, kernel_session):
+    def _start_session(self, kernel_session: dict) -> bool:
         # Attempt to start kernel from persisted state.  if started, record kernel_session in dictionary
         # else delete session
         kernel_id = kernel_session["kernel_id"]
@@ -189,7 +192,7 @@ def _start_session(self, kernel_session):
 
         return True
 
-    def delete_session(self, kernel_id):
+    def delete_session(self, kernel_id: str) -> None:
         """
         Removes saved session associated with kernel_id from dictionary and persisted storage.
         """
@@ -198,7 +201,7 @@ def delete_session(self, kernel_id):
         if self.enable_persistence:
             self.log.info("Deleted persisted kernel session for id: %s" % kernel_id)
 
-    def _delete_sessions(self, kernel_ids):
+    def _delete_sessions(self, kernel_ids: list[str]) -> None:
         # Remove unstarted sessions and rewrite
         kernels_lock.acquire()
         try:
@@ -219,7 +222,7 @@ def _delete_sessions(self, kernel_ids):
             kernels_lock.release()
 
     @staticmethod
-    def pre_save_transformation(session):
+    def pre_save_transformation(session: dict) -> dict:
         kernel_id = list(session.keys())[0]
         session_info = session[kernel_id]
         if session_info.get("connection_info"):
@@ -231,7 +234,7 @@ def pre_save_transformation(session):
         return session
 
     @staticmethod
-    def post_load_transformation(session):
+    def post_load_transformation(session: dict) -> dict:
         kernel_id = list(session.keys())[0]
         session_info = session[kernel_id]
         if session_info.get("connection_info"):
@@ -243,7 +246,7 @@ def post_load_transformation(session):
         return session
 
     # abstractmethod
-    def load_sessions(self):
+    def load_sessions(self) -> None:
         """
         Load and initialize _sessions member from persistent storage.  This method is called from start_sessions().
         """
@@ -252,33 +255,31 @@ def load_sessions(self):
         )
 
     # abstractmethod
-    def load_session(self, kernel_id):
+    def load_session(self, kernel_id: str) -> None:
         """
         Load and initialize _sessions member from persistent storage for a single kernel.  This method is called from
         refresh_sessions().
         """
-        raise NotImplementedError(
-            "KernelSessionManager.load_sessions() requires an implementation!"
-        )
+        raise NotImplementedError("KernelSessionManager.load_session() requires an implementation!")
 
     # abstractmethod
-    def delete_sessions(self, kernel_ids):
+    def delete_sessions(self, kernel_ids: list[str]) -> None:
         """
         Delete the sessions in persistent storage.  Caller is responsible for synchronizing call.
         """
         raise NotImplementedError(
             "KernelSessionManager.delete_sessions(kernel_ids) requires an implementation!"
         )
 
-    def save_session(self, kernel_id):
+    def save_session(self, kernel_id: str) -> None:
         """
         Saves the sessions dictionary to persistent store.  Caller is responsible for synchronizing call.
         """
         raise NotImplementedError(
             "KernelSessionManager.save_session(kernel_id) requires an implementation!"
         )
 
-    def active_sessions(self, username):
+    def active_sessions(self, username: str) -> int:
         """
         Returns the number of active sessions for the given username.
 
@@ -296,7 +297,7 @@ def active_sessions(self, username):
         return 0
 
     @staticmethod
-    def get_kernel_username(**kwargs):
+    def get_kernel_username(**kwargs) -> str:
         """
         Returns the kernel's logical username from env dict.
 
@@ -331,23 +332,23 @@ class FileKernelSessionManager(KernelSessionManager):
 
     # Change the default to Jupyter Data Dir.
     @default("persistence_root")
-    def persistence_root_default(self):
+    def persistence_root_default(self) -> str:
         return os.getenv(self.persistence_root_env, jupyter_data_dir())
 
-    def __init__(self, kernel_manager, **kwargs):
+    def __init__(self, kernel_manager: RemoteMappingKernelManager, **kwargs):  # noqa: F821
         super().__init__(kernel_manager, **kwargs)
         if self.enable_persistence:
             self.log.info(f"Kernel session persistence location: {self._get_sessions_loc()}")
 
-    def delete_sessions(self, kernel_ids):
+    def delete_sessions(self, kernel_ids: list[str]) -> None:
         if self.enable_persistence:
             for kernel_id in kernel_ids:
                 kernel_file_name = "".join([kernel_id, ".json"])
                 kernel_session_file_path = os.path.join(self._get_sessions_loc(), kernel_file_name)
                 if os.path.exists(kernel_session_file_path):
                     os.remove(kernel_session_file_path)
 
-    def save_session(self, kernel_id):
+    def save_session(self, kernel_id: str) -> None:
         if self.enable_persistence:
             if kernel_id is not None:
                 kernel_file_name = "".join([kernel_id, ".json"])
@@ -358,7 +359,7 @@ def save_session(self, kernel_id):
                     json.dump(KernelSessionManager.pre_save_transformation(temp_session), fp)
                     fp.close()
 
-    def load_sessions(self):
+    def load_sessions(self) -> None:
         if self.enable_persistence:
             kernel_session_files = [
                 json_files
@@ -368,21 +369,21 @@ def load_sessions(self):
             for kernel_session_file in kernel_session_files:
                 self._load_session_from_file(kernel_session_file)
 
-    def load_session(self, kernel_id):
+    def load_session(self, kernel_id: str) -> None:
         if self.enable_persistence:
             if kernel_id is not None:
                 kernel_session_file = "".join([kernel_id, ".json"])
                 self._load_session_from_file(kernel_session_file)
 
-    def _load_session_from_file(self, file_name):
+    def _load_session_from_file(self, file_name: str) -> None:
         kernel_session_file_path = os.path.join(self._get_sessions_loc(), file_name)
         if os.path.exists(kernel_session_file_path):
             self.log.debug(f"Loading saved session(s) from {kernel_session_file_path}")
             with open(kernel_session_file_path) as fp:
                 self._sessions.update(KernelSessionManager.post_load_transformation(json.load(fp)))
                 fp.close()
 
-    def _get_sessions_loc(self):
+    def _get_sessions_loc(self) -> str:
         path = os.path.join(self.persistence_root, KERNEL_SESSIONS_DIR_NAME)
         if not os.path.exists(path):
             os.makedirs(path, 0o755)
@@ -406,7 +407,7 @@ class WebhookKernelSessionManager(KernelSessionManager):
     )
 
     @default("webhook_url")
-    def webhook_url_default(self):
+    def webhook_url_default(self) -> str | None:
         return os.getenv(self.webhook_url_env, None)
 
     # Webhook Username
@@ -418,7 +419,7 @@ def webhook_url_default(self):
     )
 
     @default("webhook_username")
-    def webhook_username_default(self):
+    def webhook_username_default(self) -> str | None:
         return os.getenv(self.webhook_username_env, None)
 
     # Webhook Password
@@ -430,7 +431,7 @@ def webhook_username_default(self):
     )
 
     @default("webhook_password")
-    def webhook_password_default(self):
+    def webhook_password_default(self) -> str | None:
         return os.getenv(self.webhook_password_env, None)
 
     # Auth Type
@@ -443,10 +444,10 @@ def webhook_password_default(self):
     )
 
     @default("auth_type")
-    def auth_type_default(self):
+    def auth_type_default(self) -> str | None:
         return os.getenv(self.auth_type_env, None)
 
-    def __init__(self, kernel_manager, **kwargs):
+    def __init__(self, kernel_manager: RemoteMappingKernelManager, **kwargs):  # noqa: F821
         super().__init__(kernel_manager, **kwargs)
         if self.enable_persistence:
             self.log.info("Webhook kernel session persistence activated")
@@ -464,7 +465,7 @@ def __init__(self, kernel_manager, **kwargs):
                 else:
                     self.log.error("Username and/or password aren't set")
 
-    def delete_sessions(self, kernel_ids):
+    def delete_sessions(self, kernel_ids: list[str]) -> None:
         """
         Deletes kernel sessions from database
 
@@ -476,7 +477,7 @@ def delete_sessions(self, kernel_ids):
             if response.status_code != 204:
                 self.log.error(response.raise_for_status())
 
-    def save_session(self, kernel_id):
+    def save_session(self, kernel_id: str) -> None:
         """
         Saves kernel session to database
 
@@ -494,7 +495,7 @@ def save_session(self, kernel_id):
                 if response.status_code != 204:
                     self.log.error(response.raise_for_status())
 
-    def load_sessions(self):
+    def load_sessions(self) -> None:
         """
         Loads kernel sessions from database
         """
@@ -507,7 +508,7 @@ def load_sessions(self):
             else:
                 self.log.error(response.raise_for_status())
 
-    def load_session(self, kernel_id):
+    def load_session(self, kernel_id: str) -> None:
         """
         Loads a kernel session from database
 
@@ -522,7 +523,7 @@ def load_session(self, kernel_id):
                 else:
                     self.log.error(response.raise_for_status())
 
-    def _load_session_from_response(self, kernel_session: dict):
+    def _load_session_from_response(self, kernel_session: dict) -> None:
         """
         Loads kernel session to current session