Autogenerate types with pyre and check with mypy 3.8 (#214)

* Autogenerate types with pyre and check with mypy 3.8

* Add mypy reqs

* Almost there

* Revert ignored directories

* Drop --install-types

* Bad code change

* So so so close

* All type errors resolved locally

* Reformatted with Black

* Fix tests too

* Fix tests; move config

* Fix integration tests

* Fix more integration tests

.pre-commit-config.yaml

@@ -27,6 +27,14 @@ repos:
         language: system
         types: [python]
         pass_filenames: true
+      - id: mypy
+        name: mypy-local
+        entry: poetry run mypy
+        args: ["--install-types", "--non-interactive"]
+        language: python
+        types: [python]
+        exclude: ^examples/|^truss/test.+/|model.py$
+        pass_filenames: true
   - repo: https://github.com/tcort/markdown-link-check
     rev: "v3.10.3"
     hooks:

bin/generate_base_images.py

@@ -4,9 +4,9 @@
 import subprocess
 import sys
 import tempfile
-from argparse import ArgumentParser, BooleanOptionalAction
+from argparse import ArgumentParser, BooleanOptionalAction  # type: ignore
 from pathlib import Path
-from typing import List
+from typing import List, Optional, Set
 
 from jinja2 import Environment, FileSystemLoader
 from truss.contexts.image_builder.util import (
@@ -69,7 +69,7 @@ def _build(
     use_gpu: bool = False,
     job_type: str = "server",
     push: bool = False,
-    version_tag: str = None,
+    version_tag: Optional[str] = None,
     dry_run: bool = True,
 ):
     image_name = truss_base_image_name(job_type=job_type)
@@ -125,12 +125,12 @@ def _build(
 
 
 def _build_all(
-    job_types: List[str] = None,
-    python_versions: List[str] = None,
-    live_reload_values: List[bool] = None,
-    use_gpu_values: List[bool] = None,
+    job_types: Optional[List[str]] = None,
+    python_versions: Optional[Set[str]] = None,
+    live_reload_values: Optional[List[bool]] = None,
+    use_gpu_values: Optional[List[bool]] = None,
     push: bool = False,
-    version_tag: str = None,
+    version_tag: Optional[str] = None,
     dry_run: bool = False,
 ):
     if job_types is None:
@@ -219,7 +219,7 @@ def _build_all(
     if args.python_version == "all":
         python_versions = PYTHON_VERSIONS
     else:
-        python_versions = [args.python_version]
+        python_versions = {args.python_version}
 
     if args.job_type == "all":
         job_types = ["server", "training"]

docs/reference/client.md

@@ -20,7 +20,7 @@ Get a handle to a Truss. A Truss is a build context designed to be built as a co
 
 * `TrussHandle`: A handle to the generated Truss that provides easy access to content inside.
 
-#### `init(target_directory: str, data_files: List[str] = None, requirements_file: str = None) -> truss.truss_handle.TrussHandle`
+#### `init(target_directory: str, data_files: List[str] = None, requirements_file: Optional[str] = None) -> truss.truss_handle.TrussHandle`
 
 Initialize an empty placeholder Truss. A Truss is a build context designed to be built as a container locally or uploaded into a model serving environment. This placeholder structure can be filled to represent ML models.
 
@@ -30,7 +30,7 @@ Initialize an empty placeholder Truss. A Truss is a build context designed to be
 
 #### `kill_all()`
 
-#### create(model: Any, target_directory: str = None, data_files: List[str] = None, requirements_file: str = None) -> truss.truss_handle.TrussHandle
+#### create(model: Any, target_directory: Optional[str] = None, data_files: List[str] = None, requirements_file: Optional[str] = None) -> truss.truss_handle.TrussHandle
 
 Create a Truss with the given model. A Truss is a build context designed to
 be built as a container locally or uploaded into a model serving environment.
@@ -82,7 +82,7 @@ Add a python requirement to truss model's config.
 
 Add a system package requirement to truss model's config.
 
-#### `build_docker_image(self, build_dir: pathlib.Path = None, tag: str = None)`
+#### `build_docker_image(self, build_dir: pathlib.Path = None, tag: Optional[str] = None)`
 
 Builds docker image
 
@@ -96,11 +96,11 @@ Set up a directory to build docker image from.
 
 - docker build command.
 
-#### `docker_predict(self, request: dict, build_dir: pathlib.Path = None, tag: str = None, local_port: int = 8080, detach: bool = True)`
+#### `docker_predict(self, request: dict, build_dir: pathlib.Path = None, tag: Optional[str] = None, local_port: int = 8080, detach: bool = True)`
 
 Builds docker image, runs that as a docker container and makes a prediction request to the server running on the container. Kills the container afterwards. Mostly useful for testing.
 
-#### `docker_run(self, build_dir: pathlib.Path = None, tag: str = None, local_port: int = 8080, detach=True)`
+#### `docker_run(self, build_dir: pathlib.Path = None, tag: Optional[str] = None, local_port: int = 8080, detach=True)`
 
 Builds a docker image and runs it as a container.
 

pyproject.toml

@@ -72,6 +72,7 @@ truss = 'truss.cli:cli_group'
 
 [tool.poetry.group.dev.dependencies]
 mlflow = "^1.29.0"
+mypy = "^1.0.0"
 
 [build-system]
 requires = ["poetry-core>=1.2.1"]
@@ -80,3 +81,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.isort]
 profile = "black"
 src_paths = ["isort", "test"]
+
+[tool.mypy]
+python_version = 3.8
+ignore_missing_imports = true

truss/build.py

@@ -2,7 +2,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Optional
 
 import click
 import cloudpickle
@@ -24,51 +24,55 @@
     get_gpu_memory,
 )
 
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
 
 if is_notebook_or_ipython():
     logger.setLevel(logging.INFO)
     logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
 def populate_target_directory(
-    config: TrussConfig, target_directory_path: str = None, template: str = "custom"
+    config: TrussConfig,
+    target_directory_path: Optional[str] = None,
+    template: str = "custom",
 ) -> Path:
-
+    target_directory_path_typed = None
     if target_directory_path is None:
-        target_directory_path = build_truss_target_directory(template)
+        target_directory_path_typed = build_truss_target_directory(template)
     else:
-        target_directory_path = Path(target_directory_path)
-        target_directory_path.mkdir(parents=True, exist_ok=True)
+        target_directory_path_typed = Path(target_directory_path)
+        target_directory_path_typed.mkdir(parents=True, exist_ok=True)
 
     # Create data dir
-    (target_directory_path / config.data_dir).mkdir()
+    (target_directory_path_typed / config.data_dir).mkdir()
 
     # Create bundled packages dir
-    (target_directory_path / config.bundled_packages_dir).mkdir()
+    (target_directory_path_typed / config.bundled_packages_dir).mkdir()
 
     # Create model module dir
-    model_dir = target_directory_path / config.model_module_dir
+    model_dir = target_directory_path_typed / config.model_module_dir
     template_path = TEMPLATES_DIR / template
     copy_tree_path(template_path / "model", model_dir)
 
     examples_path = template_path / DEFAULT_EXAMPLES_FILENAME
     if examples_path.exists():
-        copy_file_path(examples_path, target_directory_path / DEFAULT_EXAMPLES_FILENAME)
+        copy_file_path(
+            examples_path, target_directory_path_typed / DEFAULT_EXAMPLES_FILENAME
+        )
 
     # Write config
-    with (target_directory_path / CONFIG_FILE).open("w") as config_file:
+    with (target_directory_path_typed / CONFIG_FILE).open("w") as config_file:
         yaml.dump(config.to_dict(), config_file)
 
-    return target_directory_path
+    return target_directory_path_typed
 
 
 def create_from_model(
     model: Any,
-    target_directory: str = None,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
+    target_directory: Optional[str] = None,
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
 ) -> TrussHandle:
     """Create a Truss with the given model. A Truss is a build context designed to
     be built as a container locally or uploaded into a model serving environment.
@@ -113,10 +117,10 @@ def create_from_model(
 
 def create_from_pipeline(
     pipeline: Callable,
-    target_directory: str = None,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
+    target_directory: Optional[str] = None,
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
 ):
     """Create a Truss from a function. A Truss is a build context designed to
     be built as a container locally or uploaded into a model serving environment.
@@ -173,10 +177,10 @@ def create_from_pipeline(
 
 def create_from_mlflow_uri(
     model_uri: str,
-    target_directory: str = None,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
+    target_directory: Optional[str] = None,
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
 ):
     """Create a Truss with the given model. A Truss is a build context designed to
     be built as a container locally or uploaded into a model serving environment.
@@ -216,9 +220,9 @@ def create_from_model_with_exception_handler(*args):
 
 def init(
     target_directory: str,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
     trainable: bool = False,
 ) -> TrussHandle:
     """
@@ -272,10 +276,10 @@ def from_directory(*args, **kwargs):
 
 def create(
     model: Any,
-    target_directory: str = None,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
+    target_directory: Optional[str] = None,
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
 ) -> TrussHandle:
     # Some model objects can are callable (like Keras models)
     # so we first attempt to make Truss via a model object
@@ -301,7 +305,7 @@ def mk_truss(*args, **kwargs):
     return create(*args, **kwargs)
 
 
-def cleanup():
+def cleanup() -> None:
     """
     Cleans up .truss directory.
     """
@@ -315,10 +319,10 @@ def cleanup():
 
 def _update_truss_props(
     scaf: TrussHandle,
-    data_files: List[str] = None,
-    requirements_file: str = None,
-    bundled_packages: List[str] = None,
-):
+    data_files: Optional[List[str]] = None,
+    requirements_file: Optional[str] = None,
+    bundled_packages: Optional[List[str]] = None,
+) -> None:
     if data_files is not None:
         for data_file in data_files:
             scaf.add_data(data_file)
@@ -334,7 +338,7 @@ def _update_truss_props(
 def _populate_default_training_code(
     config: TrussConfig,
     target_directory_path: Path,
-):
+) -> None:
     """Populate default training code in a truss.
 
     Assumes target directory already exists.
@@ -347,5 +351,5 @@ def _populate_default_training_code(
     copy_tree_path(template_path / "train", truss_training_module_dir)
 
 
-def kill_all():
+def kill_all() -> None:
     kill_containers({TRUSS: True})