Make dags play nicely with runtime type checkers

.github/workflows/main.yml

@@ -32,7 +32,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.4
         with:
-          pixi-version: v0.65.0
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           frozen: true
@@ -47,7 +47,7 @@ jobs:
         run: pixi run -e ${{ matrix.environment }} tests-with-cov
       - name: Upload coverage report.
         if: runner.os == 'Linux' && matrix.environment == 'py314'
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v6
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
   run-ty:
@@ -57,7 +57,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.4
         with:
-          pixi-version: v0.65.0
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           frozen: true

.gitignore

@@ -14,21 +14,20 @@ MANIFEST
 sdist/
 wheels/
 
-# Documentation
-docs/_build/
-
 # IDE
 .idea/
 .vscode/
 
-# Jupyter
+# Jupyter / Jupyter Book
 .ipynb_checkpoints/
+_build
 
 # macOS
 .DS_Store
 
 # pixi
 .pixi/
+node_modules/
 
 # Python
 __pycache__/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule ".ai-instructions"]
+	path = .ai-instructions
+	url = https://github.com/OpenSourceEconomics/ai-instructions.git

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
       - id: check-hooks-apply
       - id: check-useless-excludes
   - repo: https://github.com/tox-dev/pyproject-fmt
-    rev: v2.18.1
+    rev: v2.21.2
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/lyz-code/yamlfix
@@ -47,11 +47,11 @@ repos:
     hooks:
       - id: yamllint
   - repo: https://github.com/python-jsonschema/check-jsonschema
-    rev: 0.37.0
+    rev: 0.37.2
     hooks:
       - id: check-github-workflows
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.6
+    rev: v0.15.12
     hooks:
       - id: ruff-check
         args:

AGENTS.md

@@ -1,4 +1,4 @@
-@https://raw.githubusercontent.com/OpenSourceEconomics/ai-instructions/make-submodule/profiles/tier-a.md
+@.ai-instructions/profiles/tier-a.md
 
 # dags
 

CHANGES.md

@@ -4,6 +4,13 @@ This is a record of all past dags releases and what went into them in reverse
 chronological order. We follow [semantic versioning](https://semver.org/) and all
 releases are available on [conda-forge](https://anaconda.org/conda-forge/dags).
 
+## Unreleased
+
+## 0.6.0
+
+- :gh:`82` Make dags wrappers play nicely with runtime type checkers
+  (:ghuser:`hmgaudecker`).
+
 ## 0.5.1
 
 - :gh:`79` Use AGENTS.md, update hooks and rules (:ghuser:`hmgaudecker`).

GEMINI.md

@@ -0,0 +1 @@
+@AGENTS.md

pyproject.toml

@@ -60,6 +60,53 @@ build.targets.wheel.sources = [ "src" ]
 metadata.allow-direct-references = true
 version.source = "vcs"
 
+[tool.pixi.dependencies]
+jupyterlab = "*"
+prek = "*"
+python = ">=3.11,<3.15"
+[tool.pixi.environments]
+docs = [ "docs", "py314" ]
+py311 = [ "py311", "tests" ]
+py312 = [ "py312", "tests" ]
+py313 = [ "py313", "tests" ]
+py314 = [ "py314", "tests", "type-checking" ]
+[tool.pixi.feature.docs.dependencies]
+jupyter-book = ">=2.0"
+mystmd = "*"
+numpy = "*"
+[tool.pixi.feature.docs.tasks]
+build-docs = { cmd = "jupyter book build --html --execute", cwd = "docs" }
+view-docs = { cmd = "jupyter book start --execute", cwd = "docs" }
+[tool.pixi.feature.py311.dependencies]
+python = "~=3.11.0"
+[tool.pixi.feature.py312.dependencies]
+python = "~=3.12.0"
+[tool.pixi.feature.py313.dependencies]
+python = "~=3.13.0"
+[tool.pixi.feature.py314.dependencies]
+python = "~=3.14.0"
+[tool.pixi.feature.tests.dependencies]
+numpy = "*"
+[tool.pixi.feature.tests.pypi-dependencies]
+beartype = "*"
+pytest = "*"
+pytest-cov = "*"
+pytest-xdist = "*"
+[tool.pixi.feature.tests.tasks]
+tests = "pytest"
+tests-with-cov = "pytest --cov-report=xml --cov=./"
+[tool.pixi.feature.type-checking.pypi-dependencies]
+ty = "*"
+types-networkx = "*"
+[tool.pixi.feature.type-checking.tasks]
+ty = "ty check"
+[tool.pixi.pypi-dependencies]
+dags = { path = ".", editable = true }
+pdbp = "*"
+[tool.pixi.workspace]
+channels = [ "conda-forge" ]
+platforms = [ "linux-64", "osx-64", "osx-arm64", "win-64" ]
+
 [tool.ruff]
 fix = true
 unsafe-fixes = false
@@ -138,76 +185,24 @@ expand_tables = [
   "tool.pixi.workspace",
 ]
 
-[tool.pytest]
-ini_options.addopts = [ "--pdbcls=pdbp:Pdb" ]
-ini_options.filterwarnings = []
-ini_options.norecursedirs = [ "docs" ]
-
 [tool.ty]
 rules.ambiguous-protocol-member = "error"
 rules.deprecated = "error"
 rules.division-by-zero = "error"
 rules.ignore-comment-unknown-rule = "error"
-rules.invalid-argument-type = "error"
+rules.ineffective-final = "error"
+rules.invalid-enum-member-annotation = "error"
 rules.invalid-ignore-comment = "error"
-rules.invalid-return-type = "error"
-rules.possibly-missing-attribute = "error"
-rules.possibly-missing-implicit-call = "error"
-rules.possibly-missing-import = "error"
-rules.possibly-unresolved-reference = "error"
+rules.invalid-legacy-positional-parameter = "error"
 rules.redundant-cast = "error"
-rules.undefined-reveal = "error"
-rules.unresolved-global = "error"
-rules.unsupported-base = "error"
+rules.unused-awaitable = "error"
 rules.unused-ignore-comment = "error"
-rules.useless-overload-body = "error"
 src.exclude = [ "docs/**/*.ipynb" ]
 
-[tool.pixi.dependencies]
-jupyterlab = "*"
-prek = "*"
-python = ">=3.11,<3.15"
-[tool.pixi.environments]
-docs = [ "docs", "py314" ]
-py311 = [ "py311", "tests" ]
-py312 = [ "py312", "tests" ]
-py313 = [ "py313", "tests" ]
-py314 = [ "py314", "tests", "type-checking" ]
-[tool.pixi.feature.docs.dependencies]
-jupyter-book = ">=2.0"
-mystmd = "*"
-numpy = "*"
-[tool.pixi.feature.docs.tasks]
-build-docs = { cmd = "jupyter book build --html --execute", cwd = "docs" }
-view-docs = { cmd = "jupyter book start --execute", cwd = "docs" }
-[tool.pixi.feature.py311.dependencies]
-python = "~=3.11.0"
-[tool.pixi.feature.py312.dependencies]
-python = "~=3.12.0"
-[tool.pixi.feature.py313.dependencies]
-python = "~=3.13.0"
-[tool.pixi.feature.py314.dependencies]
-python = "~=3.14.0"
-[tool.pixi.feature.tests.dependencies]
-numpy = "*"
-[tool.pixi.feature.tests.pypi-dependencies]
-pytest = "*"
-pytest-cov = "*"
-pytest-xdist = "*"
-[tool.pixi.feature.tests.tasks]
-tests = "pytest"
-tests-with-cov = "pytest --cov-report=xml --cov=./"
-[tool.pixi.feature.type-checking.pypi-dependencies]
-ty = "*"
-types-networkx = "*"
-[tool.pixi.feature.type-checking.tasks]
-ty = "ty check"
-[tool.pixi.pypi-dependencies]
-dags = { path = ".", editable = true }
-pdbp = "*"
-[tool.pixi.workspace]
-channels = [ "conda-forge" ]
-platforms = [ "linux-64", "osx-64", "osx-arm64", "win-64" ]
+[tool.pytest]
+ini_options.addopts = [ "--pdbcls=pdbp:Pdb" ]
+ini_options.filterwarnings = []
+ini_options.norecursedirs = [ "docs" ]
 
 [tool.yamlfix]
 line_length = 88

src/dags/annotations.py

@@ -137,23 +137,37 @@ def _get_annotations_from_signature(
 ) -> dict[str, Any]:
     """Extract annotations from the function signature.
 
-    This is a fallback for when inspect.get_annotations returns incorrect results,
-    such as in Python 3.14's args/kwargs annotation mismatch case.
+    This is a fallback for when inspect.get_annotations returns incorrect results:
+    the Python 3.14 args/kwargs annotation mismatch, and dags wrappers
+    (`with_signature`, `rename_arguments`, the `*_output` converters), which
+    advertise the `*args, **kwargs` forwarder shape on `__annotations__` and
+    keep the user-described view on `__signature__`.
 
     """
     sig = inspect.signature(func)
     annotations: dict[str, Any] = {}
     for param_name, param in sig.parameters.items():
         if param.annotation != inspect.Parameter.empty:
-            annotations[param_name] = (
-                param.annotation
-                if eval_str or isinstance(param.annotation, str)
-                else _get_str_repr(param.annotation)
+            annotations[param_name] = _normalise_annotation(
+                param.annotation, eval_str=eval_str
             )
     if sig.return_annotation != inspect.Signature.empty:
-        annotations["return"] = (
-            sig.return_annotation
-            if eval_str or isinstance(sig.return_annotation, str)
-            else _get_str_repr(sig.return_annotation)
+        annotations["return"] = _normalise_annotation(
+            sig.return_annotation, eval_str=eval_str
         )
     return annotations
+
+
+def _normalise_annotation(annotation: Any, *, eval_str: bool) -> Any:
+    """Normalise a single signature annotation for `get_annotations`.
+
+    With `eval_str=True` the annotation is returned untouched. With
+    `eval_str=False` the caller wants strings: bare `type` objects are
+    reduced to their `__name__`, but strings and the structured return
+    annotations produced by the `*_output` converters (a `tuple` / `list`
+    / `dict` of type strings) are passed through as-is — stringifying
+    those would lose their structure.
+    """
+    if eval_str or isinstance(annotation, (str, tuple, list, dict)):
+        return annotation
+    return _get_str_repr(annotation)

src/dags/output.py

@@ -3,9 +3,11 @@
 import functools
 import inspect
 from collections.abc import Callable, Sequence
-from typing import Any, Unpack, overload
+from typing import Any, Unpack, cast, overload
 
+from dags.annotations import get_annotations
 from dags.exceptions import DagsError
+from dags.signature import forwarder_annotations
 from dags.typing import MixedTupleType, P, T
 
 
@@ -16,15 +18,24 @@ def _apply_return_annotation(
 ) -> None:
     """Apply a new return annotation to a wrapper function.
 
-    Updates both __signature__ and __annotations__ on the wrapper.
+    The user-described view (parameters from `func`, plus the new
+    `return_annotation`) is written to the wrapper's `__signature__`. The
+    wrapper's `__annotations__` advertises the `*args, **kwargs` forwarder
+    shape — the wrapper accepts anything at the Python level and only
+    `func` expects the typed arguments. See `forwarder_annotations` for
+    the rationale; `dags.get_annotations` recovers the user view from
+    `__signature__`.
+
+    `func` is itself typically a dags wrapper (e.g. the `with_signature`
+    output of `concatenate_functions`), so its return annotation lives on
+    `__signature__`, not `__annotations__` — read it via
+    `dags.get_annotations`.
     """
     signature = inspect.signature(func)
-    annotations = inspect.get_annotations(func)
-    if "return" in annotations:
+    if "return" in get_annotations(func):
         signature = signature.replace(return_annotation=return_annotation)
-        annotations["return"] = return_annotation
     wrapper.__signature__ = signature  # ty: ignore[unresolved-attribute]
-    wrapper.__annotations__ = annotations
+    wrapper.__annotations__ = forwarder_annotations()
 
 
 def single_output(
@@ -40,9 +51,13 @@ def wrapper_single_output(*args: P.args, **kwargs: P.kwargs) -> T:
         return raw[0]
 
     if set_annotations:
-        annotations = inspect.get_annotations(func)
+        annotations = get_annotations(func)
         if "return" in annotations:
-            tuple_of_types: tuple[str, ...] = annotations["return"]
+            # `func` is a tuple-returning concatenated function; its return
+            # annotation (recovered from `__signature__`) is a tuple of type
+            # strings, which `get_annotations`' `dict[str, str]` type cannot
+            # express.
+            tuple_of_types = cast("tuple[str, ...]", annotations["return"])
             _apply_return_annotation(wrapper_single_output, func, tuple_of_types[0])
 
     return wrapper_single_output
@@ -88,9 +103,9 @@ def wrapper_dict_output(*args: P.args, **kwargs: P.kwargs) -> dict[str, T]:
             return dict(zip(keys, raw, strict=True))
 
         if set_annotations:
-            annotations = inspect.get_annotations(func)
+            annotations = get_annotations(func)
             if "return" in annotations:
-                tuple_of_types: tuple[str, ...] = annotations["return"]
+                tuple_of_types = cast("tuple[str, ...]", annotations["return"])
                 return_annotation = dict(zip(keys, tuple_of_types, strict=True))
                 _apply_return_annotation(wrapper_dict_output, func, return_annotation)
 
@@ -112,9 +127,9 @@ def wrapper_list_output(*args: P.args, **kwargs: P.kwargs) -> list[T]:
         return list(raw)
 
     if set_annotations:
-        annotations = inspect.get_annotations(func)
+        annotations = get_annotations(func)
         if "return" in annotations:
-            tuple_of_types: tuple[str, ...] = annotations["return"]
+            tuple_of_types = cast("tuple[str, ...]", annotations["return"])
             _apply_return_annotation(wrapper_list_output, func, list(tuple_of_types))
 
     return wrapper_list_output