Make dags play nicely with runtime type checkers

[hmgaudecker]

Coordinated change

Part of a three-repo change adopting runtime type checking in dags-based projects:

• OpenSourceEconomics/dags#82 (this PR) — makes dags wrappers honest forwarders so runtime checkers aren't misled. Keystone; releases as dags 0.6.
• ttsim-dev/ttsim#99 — adapts ttsim's wrapper-annotation reading to the new shape. Independent, forward-compatible, mergeable anytime.
• OpenSourceEconomics/pylcm#357 — adopts the beartype claw on construction-time subpackages. Consumer; blocked on dags 0.6.

Merge order: ttsim#99 (anytime) · dags#82 → release 0.6 · pylcm#357 (after 0.6).

Why

A dags-based project has type information that only fully exists at runtime: concatenate_functions makes function composition data, so the composed callable's real signature is (*args, **kwargs), assembled at execution time. The recurring job across these PRs is to make every type-system-facing surface tell the truth about what's actually happening — and dags' wrappers weren't.

with_signature, rename_arguments, and the *_output converters all produce def wrapper(*args, **kwargs) forwarders that accept anything — yet they advertised the user's per-parameter annotations on __annotations__. A wrapper carries two views of itself: __signature__ (named params + types — for IDEs, inspect.signature, dags' own DAG resolution) and __annotations__ (for typing.get_type_hints, beartype, typeguard). Putting the user view on __annotations__ lies to runtime checkers: under beartype's import claw the wrapper gets @beartype, and beartype then enforces the wrapped function's annotations (e.g. age: float) against the forwarder's actual arguments — JAX tracers, mocks, anything.

What changes

__annotations__ now tells the truth — {"args": object, "kwargs": object}, a permissive forwarder — while the user-described view (parameter names, type strings, return annotation) stays on __signature__. dags.get_annotations recovers the user view from __signature__ via its existing args/kwargs-mismatch fallback, so dags' own machinery is unaffected.

Transparent by design

The wrapper is now deliberately invisible to runtime checkers — and that is the correct factoring, not a gap.

A natural question: shouldn't beartype check the composed callable? It can't usefully, and shouldn't. The composed callable's signature contains only the DAG's root arguments and final return — and those are already checked, because each root argument flows into a leaf node that consumes it, and the final return comes from a leaf node. The internal wiring (node A → node B) isn't in the composed signature at all; it's checked at each node's input boundary. Runtime checking belongs on the leaf node functions, where the types actually live. A checked wrapper would only duplicate the per-node checks while adding no wiring coverage — and making it work at all would require codegen'd real signatures plus resolving string forward-refs against a synthetic namespace. The honest forwarder shape keeps the wrapper out of the way so it can't inject false signals into the checking the leaf nodes already do properly.

tests/test_runtime_type_checkers.py pins this contract: it @beartype-decorates every wrapper type and the genuine concatenate_functions output (all return_types × set_annotations) and asserts beartype does not enforce the user signature against the forwarder.

Breaking change

Callers that read func.__annotations__ directly for the user-typed view must switch to inspect.signature(func) or dags.get_annotations(func). There is no opt-out — the forwarder shape is the only behaviour. (An earlier iteration of this PR added a forwarder flag; it was dropped — the forwarder shape is structurally the right default, so it is the only one.) See the CHANGES.md entry. Known in-house impact: ttsim reads __annotations__ directly in a handful of places, handled by the companion ttsim#99.

Details

• signature.py: with_signature / rename_arguments write forwarder_annotations() to __annotations__; the now-dead _create_annotations helper is removed.
• output.py: _apply_return_annotation and the four *_output converters write forwarder_annotations() to __annotations__, keep the transformed user view on __signature__, and read the inner function's return annotation via dags.get_annotations.
• annotations.py: the args/kwargs-mismatch fallback passes tuple / list / dict annotations through unchanged instead of stringifying them, so the structured return annotations produced by the *_output converters round-trip.
• tests/test_runtime_type_checkers.py: new beartype regression suite (adds beartype as a test dependency).
• Tooling: prek autoupdate and a ty version bump folded in.

Verification

• pixi run -e py314 tests → 194 passed
• pixi run ty → clean
• prek run --all-files → clean

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 96.47887% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.56%. Comparing base (36ef046) to head (cf59c04).

Files with missing lines	Patch %	Lines
tests/test_signature.py	90.19%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #82      +/-   ##
==========================================
- Coverage   96.70%   96.56%   -0.15%     
==========================================
  Files          26       27       +1     
  Lines        1306     1398      +92     
==========================================
+ Hits         1263     1350      +87     
- Misses         43       48       +5     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.