feat(tidy3d): FXC-3573 static type checking in python client (Stage 1)

.pre-commit-config.yaml

@@ -24,4 +24,11 @@ repos:
     hooks:
       - id: zizmor
         stages: [pre-commit]
-
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.13.0
+    hooks:
+      - id: mypy
+        name: mypy (type signatures)
+        files: ^tidy3d/web
+        args:
+          - --config-file=pyproject.toml

pyproject.toml

@@ -68,6 +68,7 @@ pytest-env = "^1.1.5"
 tox = { version = "*", optional = true }
 diff-cover = { version = "*", optional = true }
 zizmor = { version = "*", optional = true }
+mypy = { version = "1.13.0", optional = true }
 
 # gdstk
 gdstk = { version = ">=0.9.49", optional = true }
@@ -136,6 +137,7 @@ dev = [
   'jupyter',
   'myst-parser',
   'memory_profiler',
+  'mypy',
   'psutil',
   'nbconvert',
   'nbdime',
@@ -312,3 +314,81 @@ norecursedirs = [
 filterwarnings = "ignore::DeprecationWarning"
 testpaths = ["tidy3d", "tests", "docs"]
 python_files = "*.py"
+
+[tool.mypy]
+python_version = "3.10"
+ignore_missing_imports = true
+follow_imports = "skip"
+disallow_untyped_defs = true
+disable_error_code = [
+  "abstract",
+  "annotation-unchecked",
+  "arg-type",
+  "assert-type",
+  "assignment",
+  "attr-defined",
+  "await-not-async",
+  "call-arg",
+  "call-overload",
+  "comparison-overlap",
+  "dict-item",
+  "empty-body",
+  "exit-return",
+  "explicit-override",
+  "func-returns-value",
+  "has-type",
+  "ignore-without-code",
+  "import",
+  "import-not-found",
+  "import-untyped",
+  "index",
+  "list-item",
+  "literal-required",
+  "method-assign",
+  "misc",
+  "mutable-override",
+  "name-defined",
+  "name-match",
+  "narrowed-type-not-subtype",
+  "no-any-return",
+  "no-any-unimported",
+  "no-overload-impl",
+  "no-redef",
+  "no-untyped-call",
+  "operator",
+  "overload-cannot-match",
+  "overload-overlap",
+  "override",
+  "possibly-undefined",
+  "prop-decorator",
+  "redundant-cast",
+  "redundant-expr",
+  "redundant-self",
+  "return",
+  "return-value",
+  "safe-super",
+  "str-bytes-safe",
+  "str-format",
+  "syntax",
+  "top-level-await",
+  "truthy-bool",
+  "truthy-function",
+  "truthy-iterable",
+  "type-abstract",
+  "type-arg",
+  "type-var",
+  "typeddict-item",
+  "typeddict-readonly-mutated",
+  "typeddict-unknown-key",
+  "unimported-reveal",
+  "union-attr",
+  "unreachable",
+  "unused-awaitable",
+  "unused-coroutine",
+  "unused-ignore",
+  "used-before-def",
+  "valid-newtype",
+  "valid-type",
+  "var-annotated",
+]
+enable_error_code = ["no-untyped-def"]

tidy3d/__main__.py

@@ -9,7 +9,7 @@
 from tidy3d.web import Job
 
 
-def main(args):
+def main(args) -> None:
     """Parse args and run the corresponding tidy3d simulaton."""
 
     parser = argparse.ArgumentParser(description="Tidy3D")

tidy3d/components/autograd/derivative_utils.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field, replace
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import xarray as xr
@@ -25,12 +25,12 @@
 class LazyInterpolator:
     """Lazy wrapper for interpolators that creates them on first access."""
 
-    def __init__(self, creator_func: Callable):
+    def __init__(self, creator_func: Callable) -> None:
         """Initialize with a function that creates the interpolator when called."""
         self.creator_func = creator_func
         self._interpolator = None
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any):
         """Create interpolator on first call and delegate to it."""
         if self._interpolator is None:
             self._interpolator = self.creator_func()
@@ -172,7 +172,7 @@ class DerivativeInfo:
     # private cache for interpolators
     _interpolators_cache: dict = field(default_factory=dict, init=False, repr=False)
 
-    def updated_copy(self, **kwargs):
+    def updated_copy(self, **kwargs: Any):
         """Create a copy with updated fields."""
         kwargs.pop("deep", None)
         kwargs.pop("validate", None)
@@ -251,7 +251,7 @@ def create_interpolators(self, dtype: Optional[np.dtype] = None) -> dict:
         interpolators = {}
         coord_cache = {}
 
-        def _make_lazy_interpolator_group(field_data_dict, group_key, is_field_group=True):
+        def _make_lazy_interpolator_group(field_data_dict, group_key, is_field_group=True) -> None:
             """Helper to create a group of lazy interpolators."""
             if is_field_group:
                 interpolators[group_key] = {}

tidy3d/components/autograd/functions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import itertools
+from typing import Any
 
 import autograd.numpy as anp
 import numpy as np
@@ -98,7 +99,7 @@ def interpn(
     xi: tuple[NDArray[np.float64], ...],
     *,
     method: InterpolationType = "linear",
-    **kwargs,
+    **kwargs: Any,
 ) -> NDArray[np.float64]:
     """Interpolate over a rectilinear grid in arbitrary dimensions.
 

tidy3d/components/base.py

@@ -137,7 +137,7 @@ def skip_if_fields_missing(fields: list[str], root=False):
 
     def actual_decorator(validator):
         @wraps(validator)
-        def _validator(cls, *args, **kwargs):
+        def _validator(cls, *args: Any, **kwargs: Any):
             """New validator function."""
             values = kwargs.get("values")
             if values is None:
@@ -180,7 +180,7 @@ def _hash_self(self) -> str:
         self.to_hdf5(bf)
         return hashlib.md5(bf.getvalue()).hexdigest()
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any) -> None:
         """Init method, includes post-init validators."""
         log.begin_capture()
         super().__init__(**kwargs)
@@ -274,7 +274,7 @@ def _default(o):
 
         return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
 
-    def copy(self, deep: bool = True, validate: bool = True, **kwargs) -> Self:
+    def copy(self, deep: bool = True, validate: bool = True, **kwargs: Any) -> Self:
         """Copy a Tidy3dBaseModel.  With ``deep=True`` and ``validate=True`` as default."""
         kwargs.update(deep=deep)
         new_copy = pydantic.BaseModel.copy(self, **kwargs)
@@ -286,7 +286,7 @@ def copy(self, deep: bool = True, validate: bool = True, **kwargs) -> Self:
         return new_copy
 
     def updated_copy(
-        self, path: Optional[str] = None, deep: bool = True, validate: bool = True, **kwargs
+        self, path: Optional[str] = None, deep: bool = True, validate: bool = True, **kwargs: Any
     ) -> Self:
         """Make copy of a component instance with ``**kwargs`` indicating updated field values.
 
@@ -345,7 +345,7 @@ def updated_copy(
 
         return self._updated_copy(deep=deep, validate=validate, **{field_name: new_component})
 
-    def _updated_copy(self, deep: bool = True, validate: bool = True, **kwargs) -> Self:
+    def _updated_copy(self, deep: bool = True, validate: bool = True, **kwargs: Any) -> Self:
         """Make copy of a component instance with ``**kwargs`` indicating updated field values."""
         return self.copy(update=kwargs, deep=deep, validate=validate)
 
@@ -370,7 +370,7 @@ def from_file(
         group_path: Optional[str] = None,
         lazy: bool = False,
         on_load: Optional[Callable] = None,
-        **parse_obj_kwargs,
+        **parse_obj_kwargs: Any,
     ) -> Self:
         """Loads a :class:`Tidy3dBaseModel` from .yaml, .json, .hdf5, or .hdf5.gz file.
 
@@ -471,7 +471,7 @@ def to_file(self, fname: PathLike) -> None:
         return converter(fname=fname)
 
     @classmethod
-    def from_json(cls, fname: PathLike, **parse_obj_kwargs) -> Self:
+    def from_json(cls, fname: PathLike, **parse_obj_kwargs: Any) -> Self:
         """Load a :class:`Tidy3dBaseModel` from .json file.
 
         Parameters
@@ -536,7 +536,7 @@ def to_json(self, fname: PathLike) -> None:
             file_handle.write(json_string)
 
     @classmethod
-    def from_yaml(cls, fname: PathLike, **parse_obj_kwargs) -> Self:
+    def from_yaml(cls, fname: PathLike, **parse_obj_kwargs: Any) -> Self:
         """Loads :class:`Tidy3dBaseModel` from .yaml file.
 
         Parameters
@@ -759,7 +759,7 @@ def from_hdf5(
         fname: PathLike,
         group_path: str = "",
         custom_decoders: Optional[list[Callable]] = None,
-        **parse_obj_kwargs,
+        **parse_obj_kwargs: Any,
     ) -> Self:
         """Loads :class:`Tidy3dBaseModel` instance to .hdf5 file.
 
@@ -903,7 +903,7 @@ def from_hdf5_gz(
         fname: PathLike,
         group_path: str = "",
         custom_decoders: Optional[list[Callable]] = None,
-        **parse_obj_kwargs,
+        **parse_obj_kwargs: Any,
     ) -> Self:
         """Loads :class:`Tidy3dBaseModel` instance to .hdf5.gz file.
 
@@ -1036,7 +1036,7 @@ def _json_string(self) -> str:
         """
         return self._json()
 
-    def _json(self, indent=INDENT, exclude_unset=False, **kwargs) -> str:
+    def _json(self, indent=INDENT, exclude_unset=False, **kwargs: Any) -> str:
         """Overwrites the model ``json`` representation with some extra customized handling.
 
         Parameters
@@ -1114,7 +1114,7 @@ def _insert_traced_fields(self, field_mapping: AutogradFieldMap) -> Self:
 
         self_dict = self.dict()
 
-        def insert_value(x, path: tuple[str, ...], sub_dict: dict):
+        def insert_value(x, path: tuple[str, ...], sub_dict: dict) -> None:
             """Insert a value into the path into a dictionary."""
             current_dict = sub_dict
             for key in path[:-1]:
@@ -1349,13 +1349,13 @@ def __init__(
             self,
             fname: PathLike,
             group_path: Optional[str],
-            parse_obj_kwargs: Optional[dict[str, Any]],
+            parse_obj_kwargs: Any,
         ):
             object.__setattr__(self, "_lazy_fname", Path(fname))
             object.__setattr__(self, "_lazy_group_path", group_path)
             object.__setattr__(self, "_lazy_parse_obj_kwargs", dict(parse_obj_kwargs or {}))
 
-        def copy(self, **kwargs):
+        def copy(self, **kwargs: Any):
             """Return another lazy proxy instead of materializing."""
             return _LazyProxy(
                 self._lazy_fname,

tidy3d/components/base_sim/simulation.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Any, Optional
 
 import autograd.numpy as anp
 import pydantic.v1 as pd
@@ -251,7 +251,7 @@ def plot(
         hlim: Optional[tuple[float, float]] = None,
         vlim: Optional[tuple[float, float]] = None,
         fill_structures: bool = True,
-        **patch_kwargs,
+        **patch_kwargs: Any,
     ) -> Ax:
         """Plot each of simulation's components on a plane defined by one nonzero x,y,z coordinate.
 
@@ -482,7 +482,7 @@ def plot_boundaries(
         y: Optional[float] = None,
         z: Optional[float] = None,
         ax: Ax = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> Ax:
         """Plot the simulation boundary conditions as lines on a plane
            defined by one nonzero x,y,z coordinate.
@@ -685,7 +685,7 @@ def plot_structures_heat_conductivity(
         )
 
     @classmethod
-    def from_scene(cls, scene: Scene, **kwargs) -> AbstractSimulation:
+    def from_scene(cls, scene: Scene, **kwargs: Any) -> AbstractSimulation:
         """Create a simulation from a :class:`.Scene` instance. Must provide additional parameters
         to define a valid simulation (for example, ``size``, ``run_time``, ``grid_spec``, etc).
 

tidy3d/components/boundary.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pydantic.v1 as pd
@@ -416,7 +416,7 @@ def plot(
         y: Optional[float] = None,
         z: Optional[float] = None,
         ax: Ax = None,
-        **patch_kwargs,
+        **patch_kwargs: Any,
     ) -> Ax:
         """Plot this absorber."""
 

tidy3d/components/data/data_array.py

@@ -77,7 +77,7 @@ class DataArray(xr.DataArray):
     # stores a dictionary of attributes corresponding to the data values
     _data_attrs: dict[str, str] = {}
 
-    def __init__(self, data, *args, **kwargs):
+    def __init__(self, data, *args: Any, **kwargs: Any) -> None:
         # if data is a vanilla autograd box, convert to our box
         if isbox(data) and not is_tidy_box(data):
             data = TidyArrayBox.from_arraybox(data)
@@ -155,7 +155,7 @@ def assign_coord_attrs(cls, val):
         return val
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
+    def __modify_schema__(cls, field_schema) -> None:
         """Sets the schema of DataArray object."""
 
         schema = {
@@ -435,7 +435,7 @@ def _ag_interp(
         return self._from_temp_dataset(ds)
 
     @staticmethod
-    def _ag_interp_func(var, indexes_coords, method, **kwargs):
+    def _ag_interp_func(var, indexes_coords, method, **kwargs: Any):
         """
         Interpolate the variable `var` along the coordinates specified in `indexes_coords` using the given `method`.
 

tidy3d/components/data/dataset.py

@@ -619,7 +619,7 @@ def dispersion(self) -> ModeDispersionDataArray:
             )
         return self.dispersion_raw
 
-    def plot_field(self, *args, **kwargs):
+    def plot_field(self, *args: Any, **kwargs: Any) -> None:
         """Warn user to use the :class:`.ModeSolver` ``plot_field`` function now."""
         raise DeprecationWarning(
             "The 'plot_field()' method was moved to the 'ModeSolver' object."