|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import ast |
| 4 | +import inspect |
| 5 | +import sys |
| 6 | +from inspect import get_annotations as _get_annotations_from_inspect |
| 7 | +from typing import TYPE_CHECKING |
| 8 | +from typing import Any |
| 9 | +from typing import cast |
| 10 | + |
| 11 | +if TYPE_CHECKING: |
| 12 | + from collections.abc import Callable |
| 13 | + |
3 | 14 | __all__ = ["get_annotations"] |
4 | 15 |
|
5 | 16 |
|
6 | | -from inspect import get_annotations |
| 17 | +def get_annotations( |
| 18 | + obj: Callable[..., Any], |
| 19 | + *, |
| 20 | + globals: dict[str, Any] | None = None, # noqa: A002 |
| 21 | + locals: dict[str, Any] | None = None, # noqa: A002 |
| 22 | + eval_str: bool = False, |
| 23 | +) -> dict[str, Any]: |
| 24 | + """Return evaluated annotations with better support for deferred evaluation. |
| 25 | +
|
| 26 | + Context |
| 27 | + ------- |
| 28 | + * PEP 649 introduces deferred annotations which are only evaluated when explicitly |
| 29 | + requested. See https://peps.python.org/pep-0649/ for background and why locals can |
| 30 | + disappear between definition and evaluation time. |
| 31 | + * Python 3.14 ships :mod:`annotationlib` which exposes the raw annotation source and |
| 32 | + provides the building blocks we reuse here. The module doc explains the available |
| 33 | + formats: https://docs.python.org/3/library/annotationlib.html |
| 34 | + * Other projects run into the same constraints. Pydantic tracks their work in |
| 35 | + https://github.com/pydantic/pydantic/issues/12080; we might copy improvements from |
| 36 | + there once they settle on a stable strategy. |
| 37 | +
|
| 38 | + Rationale |
| 39 | + --------- |
| 40 | + When annotations refer to loop variables inside task generators, the locals that |
| 41 | + existed during decoration have vanished by the time pytask evaluates annotations |
| 42 | + while collecting tasks. Using :func:`inspect.get_annotations` would therefore yield |
| 43 | + the same product path for every repeated task. By asking :mod:`annotationlib` for |
| 44 | + string representations and re-evaluating them with reconstructed locals (globals, |
| 45 | + default arguments, and the frame locals captured via ``@task`` at decoration time) |
| 46 | + we recover the correct per-task values. The frame locals capture is essential for |
| 47 | + cases where loop variables are only referenced in annotations (not in the function |
| 48 | + body or closure). If any of these ingredients are missing—for example on Python |
| 49 | + versions without :mod:`annotationlib` - we fall back to the stdlib implementation, |
| 50 | + so behaviour on 3.10-3.13 remains unchanged. |
| 51 | + """ |
| 52 | + if not eval_str or not hasattr(obj, "__globals__"): |
| 53 | + return _get_annotations_from_inspect( |
| 54 | + obj, globals=globals, locals=locals, eval_str=eval_str |
| 55 | + ) |
| 56 | + |
| 57 | + if sys.version_info < (3, 14): |
| 58 | + raw_annotations = _get_annotations_from_inspect( |
| 59 | + obj, globals=globals, locals=locals, eval_str=False |
| 60 | + ) |
| 61 | + evaluation_globals = cast( |
| 62 | + "dict[str, Any]", obj.__globals__ if globals is None else globals |
| 63 | + ) |
| 64 | + evaluation_locals = evaluation_globals if locals is None else locals |
| 65 | + evaluated_annotations = {} |
| 66 | + for name, expression in raw_annotations.items(): |
| 67 | + evaluated_annotations[name] = _evaluate_annotation_expression( |
| 68 | + expression, evaluation_globals, evaluation_locals |
| 69 | + ) |
| 70 | + return evaluated_annotations |
| 71 | + |
| 72 | + import annotationlib # noqa: PLC0415 |
| 73 | + |
| 74 | + raw_annotations = annotationlib.get_annotations( |
| 75 | + obj, globals=globals, locals=locals, format=annotationlib.Format.STRING |
| 76 | + ) |
| 77 | + |
| 78 | + evaluation_globals = obj.__globals__ if globals is None else globals |
| 79 | + evaluation_locals = _build_evaluation_locals(obj, locals) |
| 80 | + |
| 81 | + evaluated_annotations = {} |
| 82 | + for name, expression in raw_annotations.items(): |
| 83 | + evaluated_annotations[name] = _evaluate_annotation_expression( |
| 84 | + expression, evaluation_globals, evaluation_locals |
| 85 | + ) |
| 86 | + |
| 87 | + return evaluated_annotations |
| 88 | + |
| 89 | + |
| 90 | +def _build_evaluation_locals( |
| 91 | + obj: Callable[..., Any], provided_locals: dict[str, Any] | None |
| 92 | +) -> dict[str, Any]: |
| 93 | + # Order matters: later updates override earlier ones. |
| 94 | + # Default arguments are lowest priority (fallbacks), then provided_locals, |
| 95 | + # then snapshot_locals (captured loop variables) have highest priority. |
| 96 | + evaluation_locals: dict[str, Any] = {} |
| 97 | + evaluation_locals.update(_get_default_argument_locals(obj)) |
| 98 | + if provided_locals: |
| 99 | + evaluation_locals.update(provided_locals) |
| 100 | + evaluation_locals.update(_get_snapshot_locals(obj)) |
| 101 | + return evaluation_locals |
| 102 | + |
| 103 | + |
| 104 | +def _get_snapshot_locals(obj: Callable[..., Any]) -> dict[str, Any]: |
| 105 | + metadata = getattr(obj, "pytask_meta", None) |
| 106 | + snapshot = getattr(metadata, "annotation_locals", None) |
| 107 | + return dict(snapshot) if snapshot else {} |
| 108 | + |
| 109 | + |
| 110 | +def _get_default_argument_locals(obj: Callable[..., Any]) -> dict[str, Any]: |
| 111 | + try: |
| 112 | + parameters = inspect.signature(obj).parameters.values() |
| 113 | + except (TypeError, ValueError): |
| 114 | + return {} |
| 115 | + |
| 116 | + defaults = {} |
| 117 | + for parameter in parameters: |
| 118 | + if parameter.default is not inspect.Parameter.empty: |
| 119 | + defaults[parameter.name] = parameter.default |
| 120 | + return defaults |
| 121 | + |
| 122 | + |
| 123 | +def _evaluate_annotation_expression( |
| 124 | + expression: Any, globals_: dict[str, Any] | None, locals_: dict[str, Any] |
| 125 | +) -> Any: |
| 126 | + if not isinstance(expression, str): |
| 127 | + return expression |
| 128 | + evaluation_globals = globals_ if globals_ is not None else {} |
| 129 | + evaluated = eval(expression, evaluation_globals, locals_) # noqa: S307 |
| 130 | + if isinstance(evaluated, str): |
| 131 | + try: |
| 132 | + literal = ast.literal_eval(expression) |
| 133 | + except (SyntaxError, ValueError): |
| 134 | + return evaluated |
| 135 | + if isinstance(literal, str): |
| 136 | + try: |
| 137 | + return eval(literal, evaluation_globals, locals_) # noqa: S307 |
| 138 | + except Exception: # noqa: BLE001 |
| 139 | + return evaluated |
| 140 | + return evaluated |
0 commit comments