more typing

Author: Carreau

jupyter_server/services/sessions/sessionmanager.py

@@ -4,7 +4,10 @@
 import os
 import pathlib
 import uuid
-from typing import Dict, List
+from typing import Any, Dict, List, NewType, Optional
+
+KernelName = NewType("KernelName", str)
+ModelName = NewType("ModelName", str)
 
 try:
     import sqlite3
@@ -251,11 +254,16 @@ def new_session_id(self) -> str:
         return str(uuid.uuid4())
 
     async def create_session(
-        self, path=None, name=None, type=None, kernel_name=None, kernel_id=None
-    ):
+        self,
+        path: Optional[str] = None,
+        name: Optional[ModelName] = None,
+        type: Optional[str] = None,
+        kernel_name: Optional[KernelName] = None,
+        kernel_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
         """Creates a session and returns its model
 
-        name: str
+        name: ModelName(str)
             Usually the model name, like the filename associated with current
             kernel.
 
@@ -278,7 +286,9 @@ async def create_session(
         self._pending_sessions.remove(record)
         return result
 
-    async def get_kernel_env(self, path: str, name: Optional[str] = None) -> Dict[str, str]:
+    async def get_kernel_env(
+        self, path: Optional[str], name: Optional[ModelName] = None
+    ) -> Dict[str, str]:
         """Return the environment variables that need to be set in the kernel
 
         path : str
@@ -296,7 +306,14 @@ async def get_kernel_env(self, path: str, name: Optional[str] = None) -> Dict[st
         assert isinstance(path, str)
         return {**os.environ, "JPY_SESSION_NAME": path}
 
-    async def start_kernel_for_session(self, session_id, path, name, type, kernel_name):
+    async def start_kernel_for_session(
+        self,
+        session_id: str,
+        path: Optional[str],
+        name: Optional[ModelName],
+        type: Optional[str],
+        kernel_name: Optional[KernelName],
+    ) -> str:
         """Start a new kernel for a given session.
 
         session_id : str

tests/services/sessions/test_manager.py

@@ -8,9 +8,11 @@
 from jupyter_server.services.contents.manager import ContentsManager
 from jupyter_server.services.kernels.kernelmanager import MappingKernelManager
 from jupyter_server.services.sessions.sessionmanager import (
+    KernelName,
     KernelSessionRecord,
     KernelSessionRecordConflict,
     KernelSessionRecordList,
+    ModelName,
     SessionManager,
 )
 
@@ -419,7 +421,7 @@ async def test_good_database_filepath(jp_runtime_dir):
     )
 
     await session_manager.create_session(
-        path="/path/to/test.ipynb", kernel_name="python", type="notebook"
+        path="/path/to/test.ipynb", kernel_name=KernelName("python"), type="notebook"
     )
     # Assert that the database file exists
     assert empty_file.exists()
@@ -451,7 +453,7 @@ async def test_session_persistence(jp_runtime_dir):
     )
 
     session = await session_manager.create_session(
-        path="/path/to/test.ipynb", kernel_name="python", type="notebook"
+        path="/path/to/test.ipynb", kernel_name=KernelName("python"), type="notebook"
     )
 
     # Assert that the database file exists
@@ -482,7 +484,7 @@ async def test_pending_kernel():
     )
     # Create a session with a slow starting kernel
     fut = session_manager.create_session(
-        path="/path/to/test.ipynb", kernel_name="python", type="notebook"
+        path="/path/to/test.ipynb", kernel_name=KernelName("python"), type="notebook"
     )
     task = asyncio.create_task(fut)
     await asyncio.sleep(0.1)
@@ -506,10 +508,10 @@ async def test_pending_kernel():
 
     # Test multiple, parallel pending kernels
     fut1 = session_manager.create_session(
-        path="/path/to/test.ipynb", kernel_name="python", type="notebook"
+        path="/path/to/test.ipynb", kernel_name=KernelName("python"), type="notebook"
     )
     fut2 = session_manager.create_session(
-        path="/path/to/test.ipynb", kernel_name="python", type="notebook"
+        path="/path/to/test.ipynb", kernel_name=KernelName("python"), type="notebook"
     )
     task1 = asyncio.create_task(fut1)
     await asyncio.sleep(0.1)