Adapt to dags' runtime-type-checker-friendly wrappers

[hmgaudecker]

Coordinated change

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

• OpenSourceEconomics/dags#82 — makes dags wrappers honest forwarders so runtime checkers aren't misled. Keystone; releases as dags 0.6.
• ttsim-dev/ttsim#99 (this PR) — 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).

Summary

dags#82 moves with_signature / rename_arguments wrappers to advertise the *args, **kwargs forwarder shape on __annotations__ (the user-described view stays on __signature__). Three sites in ttsim read __annotations__ directly off functions that may be such wrappers; this PR routes those reads through dags.get_annotations, which reads the user view regardless of where dags keeps it.

This is a pure compatibility change — ttsim does not adopt runtime type checking here.

Site	What it reads
fail_if._fail_if_root_nodes_must_be_arrays_but_are_scalars	per-argument annotations of func.function (an AggByGroupFunction/ColumnFunction, whose .function is rename_arguments-wrapped)
automatically_added_functions._convertibles	e.function.__annotations__["return"] of ParamFunctions post-remove_tree_logic
tt.vectorization._create_vectorized_annotations	inspect.get_annotations(func) where func is a ParamFunction/PolicyFunction's wrapped .function

(rename_arguments results here include those produced by dags.tree.get_one_function_without_tree_logic, which remove_tree_logic routes through.)

Why

dags.get_annotations reads the user-described view regardless of where dags keeps it:

• on the current dags release it reads __annotations__ directly (no behaviour change for ttsim);
• once dags#82 lands it transparently falls back to __signature__ via the existing args/kwargs-mismatch path.

So this PR is back-compatible with the current dags release and forward-compatible with dags#82 — it can land independently, ahead of the dags release.

Future work

ttsim's policy functions are DAG nodes in exactly the way pylcm's regime functions are. The natural next step — separate from this compatibility fix — is a scoped beartype claw analogous to pylcm#357: runtime-check the leaf node functions, where the types actually live, since the dags-composed callable is (*args, **kwargs) and a static checker can say nothing about a composition determined by data. The known precondition is the one pylcm handled: where Python scalars and JAX arrays mix, tighten internal helpers to canonical JAX types and name the Python→JAX boundary explicitly (cast at the boundary) rather than widening annotations.

Verification

• pixi run -e py314 pytest → 969 passed, 7 skipped, 2 xfailed
• pixi run -e type-checking ty → clean
• prek run --all-files → clean

Tested against the current released dags (no forwarder-shape change yet), confirming the back-compatible half. The forward-compatible half is exercised by dags#82's own suite.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[hmgaudecker]

Autoreview.