chore: Enhance type hints for potential None arguments (#2580) (#2823)

changes/2580.misc.md

@@ -0,0 +1 @@
+Enhacne type hints for potential `None` arguments

src/ai/backend/accelerator/cuda_open/plugin.py

@@ -105,7 +105,7 @@ class CUDAPlugin(AbstractComputePlugin):
     device_mask: Sequence[DeviceId] = []
     enabled: bool = True
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         rx_triple_version = re.compile(r"(\d+\.\d+\.\d+)")
 
         # Basic docker version & nvidia container runtime check

src/ai/backend/accelerator/mock/plugin.py

@@ -151,7 +151,7 @@ class MockPlugin(AbstractComputePlugin):
     nvdocker_version: Tuple[int, ...] = (0, 0, 0)
     docker_version: Tuple[int, ...] = (0, 0, 0)
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         # Read the mockup device config.
         raw_cfg, cfg_src_path = config.read_from_file(None, "mock-accelerator")
         self.mock_config = _mock_config_iv.check(raw_cfg)

src/ai/backend/agent/agent.py

@@ -312,7 +312,7 @@ def get_runner_mount(
         src: Union[str, Path],
         target: Union[str, Path],
         perm: Literal["ro", "rw"] = "ro",
-        opts: Mapping[str, Any] = None,
+        opts: Optional[Mapping[str, Any]] = None,
     ):
         """
         Return mount object to mount target krunner file/folder/volume.
@@ -788,7 +788,7 @@ async def produce_event(self, event: AbstractEvent) -> None:
 
     async def produce_error_event(
         self,
-        exc_info: Tuple[Type[BaseException], BaseException, TracebackType] = None,
+        exc_info: Optional[Tuple[Type[BaseException], BaseException, TracebackType]] = None,
     ) -> None:
         exc_type, exc, tb = sys.exc_info() if exc_info is None else exc_info
         pretty_message = "".join(traceback.format_exception_only(exc_type, exc)).strip()
@@ -1166,8 +1166,8 @@ async def inject_container_lifecycle_event(
         reason: KernelLifecycleEventReason,
         *,
         container_id: Optional[ContainerId] = None,
-        exit_code: int = None,
-        done_future: asyncio.Future = None,
+        exit_code: Optional[int] = None,
+        done_future: Optional[asyncio.Future] = None,
         suppress_events: bool = False,
     ) -> None:
         cid: Optional[ContainerId] = None

src/ai/backend/agent/docker/agent.py

@@ -430,7 +430,7 @@ def get_runner_mount(
         src: Union[str, Path],
         target: Union[str, Path],
         perm: Literal["ro", "rw"] = "ro",
-        opts: Mapping[str, Any] = None,
+        opts: Optional[Mapping[str, Any]] = None,
     ) -> Mount:
         return Mount(
             type,

src/ai/backend/agent/docker/intrinsic.py

@@ -141,7 +141,7 @@ class CPUPlugin(AbstractComputePlugin):
         (SlotName("cpu"), SlotTypes.COUNT),
     ]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:
@@ -456,7 +456,7 @@ class MemoryPlugin(AbstractComputePlugin):
         (SlotName("mem"), SlotTypes.BYTES),
     ]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:

src/ai/backend/agent/docker/utils.py

@@ -3,7 +3,7 @@
 import logging
 import subprocess
 from pathlib import Path
-from typing import Any, BinaryIO, Mapping, Tuple, cast
+from typing import Any, BinaryIO, Mapping, Optional, Tuple, cast
 
 import pkg_resources
 from aiodocker.docker import Docker
@@ -23,7 +23,7 @@ def __init__(
         image_ref: str,
         container_config: Mapping[str, Any],
         *,
-        name: str = None,
+        name: Optional[str] = None,
     ) -> None:
         self.image_ref = image_ref
         arch = get_arch_name()

src/ai/backend/agent/dummy/agent.py

@@ -153,7 +153,7 @@ def get_runner_mount(
         src: str | Path,
         target: str | Path,
         perm: Literal["ro", "rw"] = "ro",
-        opts: Mapping[str, Any] = None,
+        opts: Optional[Mapping[str, Any]] = None,
     ):
         return Mount(MountTypes.BIND, Path(), Path())
 

src/ai/backend/agent/dummy/intrinsic.py

@@ -1,6 +1,6 @@
 from decimal import Decimal
 from pathlib import Path
-from typing import Any, Collection, Mapping, Sequence
+from typing import Any, Collection, Mapping, Optional, Sequence
 
 import aiodocker
 
@@ -59,7 +59,7 @@ def __init__(
         super().__init__(plugin_config, local_config)
         self.resource_config = dummy_config["agent"]["resource"]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:
@@ -210,7 +210,7 @@ def __init__(
         super().__init__(plugin_config, local_config)
         self.resource_config = dummy_config["agent"]["resource"]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:

src/ai/backend/agent/exception.py

@@ -91,6 +91,6 @@ class AgentError(RuntimeError):
     the agent.
     """
 
-    def __init__(self, *args, exc_repr: str = None):
+    def __init__(self, *args, exc_repr: Optional[str] = None):
         super().__init__(*args)
         self.exc_repr = exc_repr

src/ai/backend/agent/kernel.py

@@ -412,7 +412,7 @@ def __init__(
         event_producer: EventProducer,
         *,
         exec_timeout: float = 0,
-        client_features: FrozenSet[str] = None,
+        client_features: Optional[FrozenSet[str]] = None,
     ) -> None:
         global _zctx
         self.kernel_id = kernel_id

src/ai/backend/agent/kubernetes/agent.py

@@ -409,7 +409,7 @@ def get_runner_mount(
         src: Union[str, Path],
         target: Union[str, Path],
         perm: Literal["ro", "rw"] = "ro",
-        opts: Mapping[str, Any] = None,
+        opts: Optional[Mapping[str, Any]] = None,
     ) -> Mount:
         return Mount(
             MountTypes.K8S_GENERIC,

src/ai/backend/agent/kubernetes/intrinsic.py

@@ -93,7 +93,7 @@ class CPUPlugin(AbstractComputePlugin):
         (SlotName("cpu"), SlotTypes.COUNT),
     ]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:
@@ -253,7 +253,7 @@ class MemoryPlugin(AbstractComputePlugin):
         (SlotName("mem"), SlotTypes.BYTES),
     ]
 
-    async def init(self, context: Any = None) -> None:
+    async def init(self, context: Optional[Any] = None) -> None:
         pass
 
     async def cleanup(self) -> None:

src/ai/backend/agent/kubernetes/utils.py

@@ -3,7 +3,7 @@
 import logging
 import subprocess
 from pathlib import Path
-from typing import Any, BinaryIO, Mapping, Tuple, cast
+from typing import Any, BinaryIO, Mapping, Optional, Tuple, cast
 
 import pkg_resources
 from aiodocker.docker import Docker
@@ -23,7 +23,7 @@ def __init__(
         image_ref: str,
         container_config: Mapping[str, Any],
         *,
-        name: str = None,
+        name: Optional[str] = None,
     ) -> None:
         self.docker = docker
         self.image_ref = image_ref

src/ai/backend/agent/resources.py

@@ -429,8 +429,8 @@ class ComputePluginContext(BasePluginContext[AbstractComputePlugin]):
     def discover_plugins(
         cls,
         plugin_group: str,
-        allowlist: set[str] = None,
-        blocklist: set[str] = None,
+        allowlist: Optional[set[str]] = None,
+        blocklist: Optional[set[str]] = None,
     ) -> Iterator[Tuple[str, Type[AbstractComputePlugin]]]:
         scanned_plugins = [*super().discover_plugins(plugin_group, allowlist, blocklist)]
 

src/ai/backend/agent/stats.py

@@ -155,7 +155,7 @@ class MovingStatistics:
     _max: Decimal
     _last: List[Tuple[Decimal, float]]
 
-    def __init__(self, initial_value: Decimal = None):
+    def __init__(self, initial_value: Optional[Decimal] = None):
         self._last = []
         if initial_value is None:
             self._sum = Decimal(0)
@@ -282,7 +282,7 @@ class StatContext:
     ]
 
     def __init__(
-        self, agent: "AbstractAgent", mode: StatModes = None, *, cache_lifespan: int = 120
+        self, agent: "AbstractAgent", mode: Optional[StatModes] = None, *, cache_lifespan: int = 120
     ) -> None:
         self.agent = agent
         self.mode = mode if mode is not None else StatModes.get_preferred_mode()

src/ai/backend/agent/utils.py

@@ -129,7 +129,7 @@ def read_sysfs(path: Union[str, Path], type_: Type[float], default: float) -> fl
 def read_sysfs(path: Union[str, Path], type_: Type[str], default: str) -> str: ...
 
 
-def read_sysfs(path: Union[str, Path], type_: Type[Any], default: Any = None) -> Any:
+def read_sysfs(path: Union[str, Path], type_: Type[Any], default: Optional[Any] = None) -> Any:
     def_vals: Mapping[Any, Any] = {
         bool: False,
         int: 0,

src/ai/backend/client/cli/admin/session.py

@@ -2,7 +2,7 @@
 
 import sys
 import uuid
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import click
 
@@ -28,7 +28,7 @@ def session() -> None:
     """
 
 
-def _list_cmd(name: str = "list", docs: str = None):
+def _list_cmd(name: str = "list", docs: Optional[str] = None):
     @pass_ctx_obj
     @click.option(
         "-s",
@@ -225,7 +225,7 @@ def list(
 session.command()(_list_cmd())
 
 
-def _info_cmd(docs: str = None):
+def _info_cmd(docs: Optional[str] = None):
     @pass_ctx_obj
     @click.argument("session_id", metavar="SESSID")
     def info(ctx: CLIContext, session_id: str) -> None:

src/ai/backend/client/cli/admin/vfolder.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sys
+from typing import Optional
 
 import click
 import humanize
@@ -24,7 +25,7 @@ def vfolder() -> None:
     """
 
 
-def _list_cmd(docs: str = None):
+def _list_cmd(docs: Optional[str] = None):
     @pass_ctx_obj
     @click.option(
         "-g",

src/ai/backend/client/cli/app.py

@@ -148,8 +148,8 @@ def __init__(
         app_name: str,
         *,
         protocol: str = "tcp",
-        args: Sequence[str] = None,
-        envs: Sequence[str] = None,
+        args: Optional[Sequence[str]] = None,
+        envs: Optional[Sequence[str]] = None,
     ) -> None:
         self.host = host
         self.port = port

src/ai/backend/client/cli/pagination.py

@@ -1,6 +1,6 @@
 import shutil
 import sys
-from typing import Any, Callable, Iterator, List, Literal, Mapping, Sequence
+from typing import Any, Callable, Iterator, List, Literal, Mapping, Optional, Sequence
 
 import click
 from tabulate import tabulate
@@ -21,8 +21,8 @@ def tabulate_items(
     items: Iterator[_Item],
     fields: Sequence[FieldSpec],
     *,
-    page_size: int = None,
-    item_formatter: Callable[[_Item], None] = None,
+    page_size: Optional[int] = None,
+    item_formatter: Optional[Callable[[_Item], None]] = None,
     tablefmt: Literal["simple", "plain", "github"] = "simple",
 ) -> Iterator[str]:
     is_first = True
@@ -76,7 +76,7 @@ def _tabulate_buffer() -> Iterator[str]:
 
 def echo_via_pager(
     text_generator: Iterator[str],
-    break_callback: Callable[[], None] = None,
+    break_callback: Optional[Callable[[], None]] = None,
 ) -> None:
     """
     A variant of ``click.echo_via_pager()`` which implements our own simplified pagination.

src/ai/backend/client/cli/session/app.py

@@ -146,8 +146,8 @@ def __init__(
         app_name: str,
         *,
         protocol: str = "tcp",
-        args: Sequence[str] = None,
-        envs: Sequence[str] = None,
+        args: Optional[Sequence[str]] = None,
+        envs: Optional[Sequence[str]] = None,
     ) -> None:
         self.host = host
         self.port = port

src/ai/backend/client/cli/session/lifecycle.py

@@ -48,7 +48,7 @@ def session():
     """Set of compute session operations"""
 
 
-def _create_cmd(docs: str = None):
+def _create_cmd(docs: Optional[str] = None):
     @click.argument("image")
     @click.option(
         "-o",
@@ -263,7 +263,7 @@ def create(
 session.command()(_create_cmd())
 
 
-def _create_from_template_cmd(docs: str = None):
+def _create_from_template_cmd(docs: Optional[str] = None):
     @click.argument("template_id")
     @click_start_option()
     @click.option(
@@ -502,7 +502,7 @@ def create_from_template(
 session.command()(_create_from_template_cmd())
 
 
-def _destroy_cmd(docs: str = None):
+def _destroy_cmd(docs: Optional[str] = None):
     @click.argument("session_names", metavar="SESSID", nargs=-1)
     @click.option(
         "-f",
@@ -572,7 +572,7 @@ def destroy(session_names, forced, owner, stats, recursive):
 session.command(aliases=["rm", "kill"])(_destroy_cmd())
 
 
-def _restart_cmd(docs: str = None):
+def _restart_cmd(docs: Optional[str] = None):
     @click.argument("session_refs", metavar="SESSION_REFS", nargs=-1)
     @click.option(
         "-o",
@@ -854,7 +854,7 @@ def abuse_history(session_id):
             sys.exit(ExitCode.FAILURE)
 
 
-def _ssh_cmd(docs: str = None):
+def _ssh_cmd(docs: Optional[str] = None):
     @click.argument("session_ref", type=str, metavar="SESSION_REF")
     @click.option(
         "-p", "--port", type=int, metavar="PORT", default=9922, help="the port number for localhost"
@@ -915,7 +915,7 @@ def ssh(ctx: click.Context, session_ref: str, port: int) -> None:
 )(_ssh_cmd())
 
 
-def _scp_cmd(docs: str = None):
+def _scp_cmd(docs: Optional[str] = None):
     @click.argument("session_ref", type=str, metavar="SESSION_REF")
     @click.argument("src", type=str, metavar="SRC")
     @click.argument("dst", type=str, metavar="DST")
@@ -1007,7 +1007,7 @@ def scp(
 )(_scp_cmd())
 
 
-def _events_cmd(docs: str = None):
+def _events_cmd(docs: Optional[str] = None):
     @click.argument("session_name_or_id", metavar="SESSION_ID_OR_NAME")
     @click.option(
         "-o",