ci: Fix up all type errors

Author: achimnol

.github/workflows/typecheck-mypy.yml

@@ -24,4 +24,4 @@ jobs:
         else
           echo "::add-matcher::.github/workflows/mypy-matcher.json"
         fi
-        python -m mypy --no-color-output src/ai/backend || exit 0
+        python -m mypy --no-color-output src/ai/backend

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

@@ -36,7 +36,7 @@ def images(operation: bool) -> None:
             print_warn('There are no registered images.')
             return
         print(tabulate((item.values() for item in items),
-                       headers=(item[0] for item in fields),
+                       headers=[item[0] for item in fields],
                        floatfmt=',.0f'))
 
 

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

@@ -128,7 +128,7 @@ class ProxyRunnerContext:
     protocol: str
     host: str
     port: int
-    args: Dict[str, str]
+    args: Dict[str, Union[None, str, List[str]]]
     envs: Dict[str, str]
     api_session: Optional[AsyncSession]
     local_server: Optional[asyncio.AbstractServer]
@@ -150,7 +150,7 @@ def __init__(self, host: str, port: int,
         self.exit_code = 0
 
         self.args, self.envs = {}, {}
-        if len(args) > 0:
+        if args is not None and len(args) > 0:
             for argline in args:
                 tokens = []
                 for token in shlex.shlex(argline,
@@ -168,7 +168,7 @@ def __init__(self, host: str, port: int,
                     self.args[tokens[0]] = tokens[1]
                 else:
                     self.args[tokens[0]] = tokens[1:]
-        if len(envs) > 0:
+        if envs is not None and len(envs) > 0:
             for envline in envs:
                 split = envline.strip().split('=', maxsplit=2)
                 if len(split) == 2:
@@ -178,6 +178,7 @@ def __init__(self, host: str, port: int,
 
     async def handle_connection(self, reader: asyncio.StreamReader,
                                 writer: asyncio.StreamWriter) -> None:
+        assert self.api_session is not None
         p = WSProxy(self.api_session, self.session_name,
                     self.app_name, self.protocol,
                     self.args, self.envs,
@@ -232,6 +233,7 @@ async def __aexit__(self, *exc_info) -> None:
             print_info("Shutting down....")
             self.local_server.close()
             await self.local_server.wait_closed()
+        assert self.api_session is not None
         await self.api_session.__aexit__(*exc_info)
         assert self.api_session.closed
         if self.local_server is not None:

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

@@ -1,6 +1,12 @@
+from __future__ import annotations
+
 import asyncio
 import json
 import re
+from typing import (
+    Union,
+    Tuple,
+)
 
 import aiohttp
 from aiohttp import web
@@ -19,6 +25,8 @@ class WebSocketProxy:
         'upstream_buffer', 'upstream_buffer_task',
     )
 
+    upstream_buffer: asyncio.Queue[Tuple[Union[str, bytes], aiohttp.WSMsgType]]
+
     def __init__(self, up_conn: aiohttp.ClientWebSocketResponse,
                  down_conn: web.WebSocketResponse):
         self.up_conn = up_conn

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

@@ -22,7 +22,7 @@
 from ..compat import asyncio_run, current_loop
 from ..exceptions import BackendError, BackendAPIError
 from ..session import Session, AsyncSession, is_legacy_server
-from ..utils import undefined
+from ..types import undefined
 from .pretty import (
     print_info, print_wait, print_done, print_error, print_fail, print_warn,
     format_info,
@@ -881,7 +881,7 @@ def start(image, name, owner,                                 # base args
               help='Set the owner of the target session explicitly.')
 # job scheduling options
 @click.option('--type', 'type_', metavar='SESSTYPE',
-              type=click.Choice(['batch', 'interactive', undefined]),
+              type=click.Choice(['batch', 'interactive', undefined]),  # type: ignore
               default=undefined,
               help='Either batch or interactive')
 @click.option('-i', '--image', default=undefined,

src/ai/backend/client/config.py

@@ -5,6 +5,8 @@
 from typing import (
     Any, Callable, Iterable, Union,
     List, Tuple, Sequence,
+    Mapping,
+    cast,
 )
 
 import appdirs
@@ -39,7 +41,7 @@ def parse_api_version(value: str) -> Tuple[int, str]:
 
 def get_env(key: str, default: Any = _undefined, *,
             clean: Callable[[str], Any] = lambda v: v):
-    '''
+    """
     Retrieves a configuration value from the environment variables.
     The given *key* is uppercased and prefixed by ``"BACKEND_"`` and then
     ``"SORNA_"`` if the former does not exist.
@@ -52,7 +54,7 @@ def get_env(key: str, default: Any = _undefined, *,
         The default is returning the value as-is.
 
     :returns: The value processed by the *clean* function.
-    '''
+    """
     key = key.upper()
     v = os.environ.get('BACKEND_' + key)
     if v is None:
@@ -73,7 +75,7 @@ def bool_env(v: str) -> bool:
     raise ValueError('Unrecognized value of boolean environment variable', v)
 
 
-def _clean_urls(v: str) -> List[URL]:
+def _clean_urls(v: Union[URL, str]) -> List[URL]:
     if isinstance(v, URL):
         return [v]
     if isinstance(v, str):
@@ -95,7 +97,7 @@ def _clean_tokens(v):
 
 
 class APIConfig:
-    '''
+    """
     Represents a set of API client configurations.
     The access key and secret key are mandatory -- they must be set in either
     environment variables or as the explicit arguments.
@@ -129,9 +131,9 @@ class APIConfig:
         access key) to be automatically mounted upon any
         :func:`Kernel.get_or_create()
         <ai.backend.client.kernel.Kernel.get_or_create>` calls.
-    '''
+    """
 
-    DEFAULTS = {
+    DEFAULTS: Mapping[str, Any] = {
         'endpoint': 'https://api.backend.ai',
         'endpoint_type': 'api',
         'version': f'v{API_VERSION[0]}.{API_VERSION[1]}',
@@ -141,9 +143,12 @@ class APIConfig:
         'connection_timeout': 10.0,
         'read_timeout': None,
     }
-    '''
+    """
     The default values except the access and secret keys.
-    '''
+    """
+
+    _group: str
+    _hash_type: str
 
     def __init__(self, *,
                  endpoint: Union[URL, str] = None,
@@ -159,17 +164,17 @@ def __init__(self, *,
                  skip_sslcert_validation: bool = None,
                  connection_timeout: float = None,
                  read_timeout: float = None) -> None:
-        from . import get_user_agent  # noqa; to avoid circular imports
+        from . import get_user_agent
         self._endpoints = (
             _clean_urls(endpoint) if endpoint else
             get_env('ENDPOINT', self.DEFAULTS['endpoint'], clean=_clean_urls))
         random.shuffle(self._endpoints)
-        self._endpoint_type = endpoint_type if endpoint_type \
+        self._endpoint_type = endpoint_type if endpoint_type is not None \
                               else get_env('ENDPOINT_TYPE', self.DEFAULTS['endpoint_type'])
-        self._domain = domain if domain else get_env('DOMAIN', self.DEFAULTS['domain'])
-        self._group = group if group else get_env('GROUP', self.DEFAULTS['group'])
-        self._version = version if version else self.DEFAULTS['version']
-        self._user_agent = user_agent if user_agent else get_user_agent()
+        self._domain = domain if domain is not None else get_env('DOMAIN', self.DEFAULTS['domain'])
+        self._group = group if group is not None else get_env('GROUP', self.DEFAULTS['group'])
+        self._version = version if version is not None else self.DEFAULTS['version']
+        self._user_agent = user_agent if user_agent is not None else get_user_agent()
         if self._endpoint_type == 'api':
             self._access_key = access_key if access_key is not None \
                                else get_env('ACCESS_KEY', '')
@@ -178,8 +183,8 @@ def __init__(self, *,
         else:
             self._access_key = 'dummy'
             self._secret_key = 'dummy'
-        self._hash_type = hash_type.lower() if hash_type else \
-                          self.DEFAULTS['hash_type']
+        self._hash_type = hash_type.lower() if hash_type is not None else \
+                          cast(str, self.DEFAULTS['hash_type'])
         arg_vfolders = set(vfolder_mounts) if vfolder_mounts else set()
         env_vfolders = set(get_env('VFOLDER_MOUNTS', [], clean=_clean_tokens))
         self._vfolder_mounts = [*(arg_vfolders | env_vfolders)]
@@ -198,16 +203,16 @@ def is_anonymous(self) -> bool:
 
     @property
     def endpoint(self) -> URL:
-        '''
+        """
         The currently active endpoint URL.
         This may change if there are multiple configured endpoints
         and the current one is not accessible.
-        '''
+        """
         return self._endpoints[0]
 
     @property
     def endpoints(self) -> Sequence[URL]:
-        '''All configured endpoint URLs.'''
+        """All configured endpoint URLs."""
         return self._endpoints
 
     def rotate_endpoints(self):
@@ -217,83 +222,83 @@ def rotate_endpoints(self):
 
     @property
     def endpoint_type(self) -> str:
-        '''
+        """
         The configured endpoint type.
-        '''
+        """
         return self._endpoint_type
 
     @property
     def domain(self) -> str:
-        '''The configured domain.'''
+        """The configured domain."""
         return self._domain
 
     @property
     def group(self) -> str:
-        '''The configured group.'''
+        """The configured group."""
         return self._group
 
     @property
     def user_agent(self) -> str:
-        '''The configured user agent string.'''
+        """The configured user agent string."""
         return self._user_agent
 
     @property
     def access_key(self) -> str:
-        '''The configured API access key.'''
+        """The configured API access key."""
         return self._access_key
 
     @property
     def secret_key(self) -> str:
-        '''The configured API secret key.'''
+        """The configured API secret key."""
         return self._secret_key
 
     @property
     def version(self) -> str:
-        '''The configured API protocol version.'''
+        """The configured API protocol version."""
         return self._version
 
     @property
     def hash_type(self) -> str:
-        '''The configured hash algorithm for API authentication signatures.'''
+        """The configured hash algorithm for API authentication signatures."""
         return self._hash_type
 
     @property
-    def vfolder_mounts(self) -> Tuple[str, ...]:
-        '''The configured auto-mounted vfolder list.'''
+    def vfolder_mounts(self) -> Sequence[str]:
+        """The configured auto-mounted vfolder list."""
         return self._vfolder_mounts
 
     @property
     def skip_sslcert_validation(self) -> bool:
-        '''Whether to skip SSL certificate validation for the API gateway.'''
+        """Whether to skip SSL certificate validation for the API gateway."""
         return self._skip_sslcert_validation
 
     @property
     def connection_timeout(self) -> float:
-        '''The maximum allowed duration for making TCP connections to the server.'''
+        """The maximum allowed duration for making TCP connections to the server."""
         return self._connection_timeout
 
     @property
     def read_timeout(self) -> float:
-        '''The maximum allowed waiting time for the first byte of the response from the server.'''
+        """The maximum allowed waiting time for the first byte of the response from the server."""
         return self._read_timeout
 
 
 def get_config():
-    '''
+    """
     Returns the configuration for the current process.
     If there is no explicitly set :class:`APIConfig` instance,
     it will generate a new one from the current environment variables
     and defaults.
-    '''
+    """
     global _config
     if _config is None:
         _config = APIConfig()
     return _config
 
 
 def set_config(conf: APIConfig):
-    '''
+    """
     Sets the configuration used throughout the current process.
-    '''
+    """
     global _config
     _config = conf

src/ai/backend/client/func/base.py

@@ -28,19 +28,19 @@ def _method(*args, **kwargs):
 
 
 def api_function(meth):
-    '''
+    """
     Mark the wrapped method as the API function method.
-    '''
+    """
     setattr(meth, '_backend_api', True)
     return meth
 
 
 class APIFunctionMeta(type):
-    '''
+    """
     Converts all methods marked with :func:`api_function` into
     session-aware methods that are either plain Python functions
     or coroutines.
-    '''
+    """
     _async = True
 
     def __init__(cls, name, bases, attrs, **kwargs):

src/ai/backend/client/func/dotfile.py

@@ -29,18 +29,16 @@ async def create(cls,
         }
         rqst.set_json(body)
         async with rqst.fetch() as resp:
-            if resp.status == 200:
-                await resp.json()
-                return cls(path, owner_access_key=owner_access_key)
+            await resp.json()
+        return cls(path, owner_access_key=owner_access_key)
 
     @api_function
     @classmethod
     async def list_dotfiles(cls) -> 'List[Mapping[str, str]]':
         rqst = Request(api_session.get(),
                        'GET', '/user-config/dotfiles')
         async with rqst.fetch() as resp:
-            if resp.status == 200:
-                return await resp.json()
+            return await resp.json()
 
     def __init__(self, path: str, owner_access_key: str = None):
         self.path = path
@@ -55,8 +53,7 @@ async def get(self) -> str:
                        'GET', f'/user-config/dotfiles',
                        params=params)
         async with rqst.fetch() as resp:
-            if resp.status == 200:
-                return await resp.json()
+            return await resp.json()
 
     @api_function
     async def update(self, data: str, permission: str):