fix type checkers (#933)

* fix type checkers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_field_type_annotations.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_field_type_annotations.py

* fix type hinting

* fix duplicate name

* typo

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

pyproject.toml

@@ -36,6 +36,7 @@ dev = [
     "ipykernel>=6.29.5",
     "mlflow>=2.20.0",
     "pre-commit>=4.1.0",
+    "pyright>=1.1.403",
     "pytest>=8.3.4",
     "pytest-benchmark[histogram]>=5.1.0",
     "pytest-cov>=6.0.0",

tests/typing/test_field_type_annotations.py

@@ -0,0 +1,136 @@
+"""Test file to verify type annotations work correctly for zntrack fields.
+
+This test should pass type checking with pyright/mypy.
+The main goal is to ensure that patterns like:
+    field: int = zntrack.params()
+    field: str = zntrack.outs()
+do not raise type errors when no explicit value is provided.
+"""
+
+from pathlib import Path
+
+import pandas as pd
+
+import zntrack
+
+
+class TestFieldAnnotations(zntrack.Node):
+    """Test basic field type annotations without explicit values."""
+
+    # These should NOT raise type errors (the main fix)
+    param_int: int = zntrack.params()
+    param_str: str = zntrack.params()
+    param_dict: dict = zntrack.params()
+
+    out_int: int = zntrack.outs()
+    out_str: str = zntrack.outs()
+    out_dict: dict = zntrack.outs()
+
+    metric_dict: dict = zntrack.metrics()
+
+    dep_int: int = zntrack.deps()
+    dep_str: str = zntrack.deps()
+
+    plots_data: pd.DataFrame = (
+        zntrack.plots()
+    )  # pandas.DataFrame, but we don't enforce the type
+
+    # Path fields
+    params_path_str: str = zntrack.params_path()
+    params_path_path: Path = zntrack.params_path()
+    outs_path_str: str = zntrack.outs_path()
+    outs_path_path: Path = zntrack.outs_path()
+    deps_path_str: str = zntrack.deps_path()
+    deps_path_path: Path = zntrack.deps_path()
+    plots_path_path: Path = zntrack.plots_path()
+    metrics_path_path: Path = zntrack.metrics_path()
+
+    def run(self):
+        pass
+
+
+class TestFieldAnnotationsWithDefaults(zntrack.Node):
+    """Test that type safety is maintained when providing explicit values."""
+
+    # These should work fine (correct type matching)
+    param_good_int: int = zntrack.params(42)
+    param_good_str: str = zntrack.params("hello")
+    param_good_tuple: str = zntrack.params(("word",))
+
+    # Using default_factory
+    param_good_dict: dict = zntrack.params(default_factory=dict)
+    param_factory: list = zntrack.params(default_factory=list)
+
+    # Path fields with correct types
+    params_path_str_with_val: str = zntrack.params_path("config.yaml")
+    params_path_path_with_val: Path = zntrack.params_path(Path("config.yaml"))
+
+    # Path fields using zntrack.nwd
+    outs_path_in_nwd: Path = zntrack.outs_path(zntrack.nwd / "output.txt")
+    metrics_path_in_nwd: Path = zntrack.metrics_path(zntrack.nwd / "metrics.json")
+    plots_path_in_nwd: Path = zntrack.plots_path(zntrack.nwd / "plot.png")
+
+    # plots_path_list
+    plots_path_list: list = zntrack.plots_path(
+        default_factory=lambda: ["plot1.png", "plot2.png"]
+    )
+
+    # path using tuples
+    params_path_tuple: tuple[Path, ...] = zntrack.params_path(
+        (zntrack.nwd / "config1.yaml", zntrack.nwd / "config2.yaml")
+    )
+    outs_path_tuple: tuple[Path, ...] = zntrack.outs_path(
+        (zntrack.nwd / "output1.txt", zntrack.nwd / "output2.txt")
+    )
+
+    def run(self):
+        pass
+
+
+class TestTypeSafetyErrors(zntrack.Node):
+    """These should cause type errors to ensure type safety is maintained."""
+
+    # # Type mismatches with explicit values (should error when uncommented)
+    param_bad_int: int = zntrack.params(
+        "string"
+    )  # Should error: str not assignable to int
+    param_bad_str: str = zntrack.params(42)  # Should error: int not assignable to str
+    param_bad_dict: dict = zntrack.params(
+        "not_a_dict"
+    )  # Should error: str not assignable to dict
+
+    # # # Path field type mismatches (should error when uncommented)
+    params_path_bad: Path = zntrack.params_path(
+        1234
+    )  # Should error: int not assignable to path types
+    outs_path_bad: str = zntrack.outs_path(
+        42
+    )  # Should error: int not assignable to path types
+    deps_path_bad: Path = zntrack.deps_path(
+        3.14
+    )  # Should error: float not assignable to path types
+
+    # List vs single type mismatches (should error when uncommented)
+    outs_path_list_mismatch: str = zntrack.outs_path(
+        ("file1.txt", "file2.txt")
+    )  # Should error: list not assignable to str
+    params_path_list_mismatch: Path = zntrack.params_path(
+        ("config1.yaml", "config2.yaml")
+    )  # Should error: list not assignable to Path
+    # using default factory with list
+    params_path_list_mismatch_lambda: Path = zntrack.params_path(
+        default_factory=lambda: ["config1.yaml", "config2.yaml"]
+    )  # Should error: list not assignable to Path
+
+    # Factory function returning wrong type (should error when uncommented)
+    param_factory_bad: int = zntrack.params(
+        default_factory=str
+    )  # Should error: str() returns str, not int
+    param_factory_bad2: dict = zntrack.params(
+        default_factory=list
+    )  # Should error: list() returns list, not dict
+
+    # zntrack.nwd as string
+    nwd_as_str: str = zntrack.outs_path(
+        zntrack.nwd / "output.txt"
+    )  # Should error: str not assignable to Path

uv.lock

@@ -12,6 +12,8 @@
 from zntrack.node import Node
 from zntrack.utils.filesystem import resolve_state_file_path
 
+_T = t.TypeVar("_T")
+
 
 def _deps_getter(self: "Node", name: str):
     zntrack_path = resolve_state_file_path(
@@ -84,6 +86,14 @@ def _deps_getter(self: "Node", name: str):
         return content
 
 
+@t.overload
+def deps() -> t.Any: ...
+
+
+@t.overload
+def deps(default: _T, **kwargs) -> _T: ...
+
+
 def deps(default=dataclasses.MISSING, **kwargs) -> t.Any:
     """Define dependencies for a node.
 

zntrack/fields/deps.py

@@ -36,6 +36,10 @@ def _metrics_save_func(self: "Node", name: str, suffix: str):
         raise TypeError(f"Error while saving {name} to {self.nwd / name}.json") from err
 
 
+@t.overload
+def outs(*, cache: bool = True, independent: bool = False, **kwargs) -> t.Any: ...
+
+
 def outs(*, cache: bool = True, independent: bool = False, **kwargs) -> t.Any:
     """Define output for a node.
 
@@ -72,6 +76,12 @@ def outs(*, cache: bool = True, independent: bool = False, **kwargs) -> t.Any:
     )
 
 
+@t.overload
+def metrics(
+    *, cache: bool | None = None, independent: bool = False, **kwargs
+) -> t.Any: ...
+
+
 def metrics(*, cache: bool | None = None, independent: bool = False, **kwargs) -> t.Any:
     """Define metrics for a node.
 

zntrack/fields/outs_and_metrics.py

@@ -21,6 +21,10 @@ def _params_getter(self: "Node", name: str):
 
 
 # Overloads for type checking
+@t.overload
+def params() -> t.Any: ...
+
+
 @t.overload
 def params(default: _T, **kwargs) -> _T: ...
 

zntrack/fields/params.py

@@ -1,3 +1,5 @@
+import typing as t
+
 import pandas as pd
 
 from zntrack.config import NOT_AVAILABLE, ZNTRACK_OPTION_PLOTS_CONFIG, FieldTypes
@@ -24,6 +26,22 @@ def _plots_getter(self: "Node", name: str, suffix: str):
         return pd.read_csv(f, index_col=0)
 
 
+@t.overload
+def plots(
+    *,
+    y: str | list[str] | None = None,
+    cache: bool = True,
+    independent: bool = False,
+    x: str = "step",
+    x_label: str | None = None,
+    y_label: str | None = None,
+    template: str | None = None,
+    title: str | None = None,
+    autosave: bool = False,
+    **kwargs,
+) -> t.Any: ...
+
+
 def plots(
     *,
     y: str | list[str] | None = None,