Use pyright for static type checking of implemenation code

.github/workflows/check.yml

@@ -15,8 +15,8 @@ jobs:
 
       - run: inv lint
 
-  mypy:
-    name: Mypy
+  type_check:
+    name: Static type checks
     runs-on: ubuntu-20.04
 
     steps:
@@ -33,7 +33,7 @@ jobs:
       - name: Build type stubs
         run: inv type-stubs
 
-      - run: inv mypy
+      - run: inv type-check
 
   check-copyright:
     name: Check copyright

modal/_mount_utils.py

@@ -1,6 +1,6 @@
 # Copyright Modal Labs 2022
-import os
 import posixpath
+import typing
 from pathlib import PurePath, PurePosixPath
 from typing import TYPE_CHECKING, Dict, List, Mapping, Tuple, Union
 
@@ -12,10 +12,13 @@
     from .s3mount import _S3Mount
 
 
+T = typing.TypeVar("T", bound=Union["_Volume", "_NetworkFileSystem", "_S3Mount"])
+
+
 def validate_mount_points(
     display_name: str,
-    volume_likes: Mapping[Union[str, os.PathLike], Union["_Volume", "_NetworkFileSystem", "_S3Mount"]],
-) -> List[Tuple[str, Union["_Volume", "_NetworkFileSystem", "_S3Mount"]]]:
+    volume_likes: Mapping[Union[str, PurePosixPath], T],
+) -> List[Tuple[str, T]]:
     """Mount point path validation for volumes and network file systems."""
 
     validated = []

modal/_output.py

@@ -237,7 +237,7 @@ def function_queueing_progress(self) -> Progress:
                 self._current_render_group.renderables.append(self._function_queueing_progress)
         return self._function_queueing_progress
 
-    def function_progress_callback(self, tag: str, total: int) -> Callable[[int, int], None]:
+    def function_progress_callback(self, tag: str, total: Optional[int]) -> Callable[[int, int], None]:
         """Adds a task to the current function_progress instance, and returns a callback
         to update task progress with new completed and total counts."""
 

modal/image.py

@@ -6,7 +6,7 @@
 import typing
 import warnings
 from inspect import isfunction
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import toml
@@ -1265,8 +1265,8 @@ def run_function(
         secrets: Sequence[_Secret] = (),  # Optional Modal Secret objects with environment variables for the container
         gpu: GPU_T = None,  # GPU specification as string ("any", "T4", "A10G", ...) or object (`modal.GPU.A100()`, ...)
         mounts: Sequence[_Mount] = (),
-        shared_volumes: Dict[Union[str, os.PathLike], _NetworkFileSystem] = {},
-        network_file_systems: Dict[Union[str, os.PathLike], _NetworkFileSystem] = {},
+        shared_volumes: Dict[Union[str, PurePosixPath], _NetworkFileSystem] = {},
+        network_file_systems: Dict[Union[str, PurePosixPath], _NetworkFileSystem] = {},
         cpu: Optional[float] = None,  # How many CPU cores to request. This is a soft limit.
         memory: Optional[int] = None,  # How much memory to request, in MiB. This is a soft limit.
         timeout: Optional[int] = 86400,  # Maximum execution time of the function in seconds.

modal/stub.py

@@ -549,7 +549,6 @@ def wrapped(
                 keep_warm=keep_warm,
                 cloud=cloud,
                 webhook_config=webhook_config,
-                cls=_cls,
                 checkpointing_enabled=checkpointing_enabled,
                 allow_background_volume_commits=_allow_background_volume_commits,
                 block_network=block_network,

modal_proto/api.proto

@@ -801,7 +801,7 @@ message Function {
 
   uint32 task_idle_timeout_secs = 25;
 
-  CloudProvider cloud_provider = 26;
+  optional CloudProvider cloud_provider = 26;
 
   uint32 warm_pool_size = 27;
 

requirements.dev.txt

@@ -28,3 +28,4 @@ wheel~=0.37.1
 nbclient==0.6.8
 notebook==6.5.1
 jupytext==1.14.1
+pyright==1.1.351

setup.cfg

@@ -19,7 +19,7 @@ packages = find:
 python_requires = >=3.8
 install_requires =
     aiohttp
-    aiostream
+    aiostream~=0.5.2
     asgiref
     certifi
     # These are pinned to be protocol-compatible with the version of

tasks.py

@@ -40,13 +40,16 @@ def lint(ctx, fix=False):
 
 
 @task
-def mypy(ctx):
-    mypy_allowlist = [
+def type_check(ctx):
+    # mypy will not check the *implementation* (.py) for files that also have .pyi type stubs
+    ctx.run("mypy . --exclude=playground --exclude=venv311 --exclude=venv38", pty=True)
+
+    # use pyright for checking implementation of those files
+    pyright_allowlist = [
         "modal/functions.py",
     ]
 
-    ctx.run("mypy .", pty=True)
-    ctx.run(f"mypy {' '.join(mypy_allowlist)} --follow-imports=skip", pty=True)
+    ctx.run(f"pyright {' '.join(pyright_allowlist)}", pty=True)
 
 
 @task