fix: preserve definition locals for postponed task annotations

.github/workflows/main.yml

@@ -11,7 +11,7 @@ on:
     - main
   pull_request:
     branches:
-    - '*'
+    - '**'
   merge_group:
 
 permissions: {}

.github/workflows/zizmor.yml

@@ -6,7 +6,7 @@ on:
     - main
   pull_request:
     branches:
-    - '*'
+    - '**'
   schedule:
     - cron: '0 7 * * 1'
   workflow_dispatch:

CHANGELOG.md

@@ -7,6 +7,10 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 
 ## Unreleased
 
+- [#1026](https://github.com/pytask-dev/pytask/pull/1026) preserves definition-local
+  bindings for generated tasks with postponed annotations and releases temporary
+  snapshots after collection or generator failures. Internal task registration no
+  longer shadows annotation globals with pytask's own locals.
 - [#1025](https://github.com/pytask-dev/pytask/pull/1025) fixes postponed annotation
   evaluation across Python versions and avoids eager evaluation during task signature
   inspection. Parameter defaults no longer override module globals or builtins when

docs/source/how_to_guides/using_task_returns.md

@@ -25,6 +25,13 @@ It works because internally the path is converted to a
 [`pytask.PathNode`](../reference_guides/api/nodes_and_tasks.md#pytask.PathNode) that is
 able to store objects of type `str` and `bytes`.
 
+When `@pytask.task` decorates a function inside another function, such as a task
+generator, pytask takes a shallow snapshot of the enclosing function's local bindings to
+evaluate annotations during collection. This also works with
+`from __future__ import annotations`, so each generated task can refer to its own loop
+values. The referenced objects are shared, and the snapshot is released after task
+collection. Module-level future annotations continue to use current module globals.
+
 !!! note
 
     Read the explanation on [nodes](../how_to_guides/writing_custom_nodes.md) to learn more

src/_pytask/collect.py

@@ -70,44 +70,54 @@ def pytask_collect(session: Session) -> bool:
     """Collect tasks."""
     session.collection_start = time.time()
 
-    _collect_from_paths(session)
-    _collect_from_tasks(session)
-    _collect_not_collected_tasks(session)
-
-    session.tasks.extend(
-        i.node
-        for i in session.collection_reports
-        if i.outcome == CollectionOutcome.SUCCESS and isinstance(i.node, PTask)
-    )
-
     try:
-        session.hook.pytask_collect_modify_tasks(session=session, tasks=session.tasks)
-        validate_unique_task_signatures(session.tasks)
-    except Exception:  # noqa: BLE001  # pragma: no cover
-        report = CollectionReport.from_exception(
-            outcome=CollectionOutcome.FAIL, exc_info=sys.exc_info()
+        _collect_from_paths(session)
+        _collect_from_tasks(session)
+        _collect_not_collected_tasks(session)
+
+        session.tasks.extend(
+            i.node
+            for i in session.collection_reports
+            if i.outcome == CollectionOutcome.SUCCESS and isinstance(i.node, PTask)
         )
-        session.collection_reports.append(report)
 
-    session.hook.pytask_collect_log(
-        session=session, reports=session.collection_reports, tasks=session.tasks
-    )
+        try:
+            session.hook.pytask_collect_modify_tasks(
+                session=session, tasks=session.tasks
+            )
+            validate_unique_task_signatures(session.tasks)
+        except Exception:  # noqa: BLE001  # pragma: no cover
+            report = CollectionReport.from_exception(
+                outcome=CollectionOutcome.FAIL, exc_info=sys.exc_info()
+            )
+            session.collection_reports.append(report)
 
-    _clear_annotation_locals(session.tasks)
+        session.hook.pytask_collect_log(
+            session=session, reports=session.collection_reports, tasks=session.tasks
+        )
 
-    return True
+        return True
+    finally:
+        _clear_annotation_locals(session.tasks)
+        _clear_annotation_locals(
+            report.node
+            for report in session.collection_reports
+            if isinstance(report.node, PTask)
+        )
 
 
-def _clear_annotation_locals(tasks: list[PTask]) -> None:
+def _clear_annotation_locals(tasks: Iterable[PTask | TaskFunction]) -> None:
     """Drop decoration-time locals snapshots once collection finishes.
 
     The snapshot is only needed to evaluate deferred annotations while collecting
-    dependencies/products. Keeping it afterwards can retain non-picklable objects (for
-    example locks) and break parallel backends that cloudpickle task functions.
+    dependencies/products and running collection hooks. Keeping it afterwards can
+    retain non-picklable objects (for example locks) and break parallel backends that
+    cloudpickle task functions.
     """
     for task in tasks:
-        if isinstance(task.function, TaskFunction):
-            task.function.pytask_meta.annotation_locals = None
+        function = task.function if isinstance(task, PTask) else task
+        if isinstance(function, TaskFunction):
+            function.pytask_meta.annotation_locals = None
 
 
 def _collect_from_paths(session: Session) -> None:
@@ -135,7 +145,11 @@ def _collect_from_tasks(session: Session) -> None:
         if is_task_function(raw_task):
             if not isinstance(raw_task, TaskFunction):
                 raw_task = task_decorator()(raw_task)  # noqa: PLW2901
+                # This registration frame is not the callable's definition scope.
+                raw_task.pytask_meta.annotation_locals = None
 
+            # Keep `path` in sync with the collision fixture in
+            # test_programmatic_annotation_globals_do_not_capture_collection_locals.
             path = get_file(raw_task)
             name = raw_task.pytask_meta.name
 

src/_pytask/models.py

@@ -42,8 +42,9 @@ class CollectionMetadata:
         A dictionary containing keyword arguments which are passed to the task when it
         is executed.
     annotation_locals
-        A snapshot of local variables captured during decoration which helps evaluate
-        deferred annotations later on.
+        A shallow snapshot of definition-local bindings captured during decoration
+        which helps evaluate deferred annotations later on. Values are shared, and
+        the snapshot is released after task collection.
     markers
         A list of markers that are attached to the task.
     name

src/_pytask/provisional.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 
+from _pytask.collect import _clear_annotation_locals
 from _pytask.config import hookimpl
 from _pytask.dag import create_dag_from_session
 from _pytask.exceptions import CollectionError
@@ -30,6 +31,7 @@
     from _pytask.reports import CollectionReport
     from _pytask.reports import ExecutionReport
     from _pytask.session import Session
+    from _pytask.typing import TaskFunction
 
 
 @hookimpl(tryfirst=True)
@@ -53,49 +55,67 @@ def pytask_execute_task(session: Session, task: PTask) -> bool | None:
     if not is_task_generator(task):
         return None
 
-    execute_task(task)
+    registered_functions = {
+        id(function) for functions in COLLECTED_TASKS.values() for function in functions
+    }
+    try:
+        execute_task(task)
+    except BaseException:
+        # Children can be registered before the generator raises, without reaching
+        # collection below. Release only their snapshots and preserve the exception.
+        _clear_annotation_locals(
+            function
+            for functions in COLLECTED_TASKS.values()
+            for function in functions
+            if id(function) not in registered_functions
+        )
+        raise
 
-    name_to_function: Mapping[str, Callable[..., Any] | PTask]
-    if isinstance(task, PTaskWithPath) and task.path in COLLECTED_TASKS:
-        tasks = COLLECTED_TASKS.pop(task.path)
-        name_to_function = parse_collected_tasks_with_task_marker(tasks)
-    elif None in COLLECTED_TASKS:
-        tasks = COLLECTED_TASKS.pop(None)
+    tasks: list[TaskFunction] = []
+    try:
+        if isinstance(task, PTaskWithPath) and task.path in COLLECTED_TASKS:
+            tasks = COLLECTED_TASKS.pop(task.path)
+        elif None in COLLECTED_TASKS:
+            tasks = COLLECTED_TASKS.pop(None)
+        else:
+            msg = f"The task generator {task.name!r} did not create any tasks."
+            raise RuntimeError(msg)
+
+        name_to_function: Mapping[str, Callable[..., Any] | PTask]
         name_to_function = parse_collected_tasks_with_task_marker(tasks)
-    else:
-        msg = f"The task generator {task.name!r} did not create any tasks."
-        raise RuntimeError(msg)
-
-    new_reports: list[CollectionReport] = []
-    for name, function in name_to_function.items():
-        report = session.hook.pytask_collect_task_protocol(
-            session=session,
-            reports=new_reports,
-            path=task.path if isinstance(task, PTaskWithPath) else None,
-            name=name,
-            obj=function,
-        )
-        if report is not None:
-            new_reports.append(report)
-
-    session.collection_reports.extend(new_reports)
-    failed_reports = [
-        report for report in new_reports if report.outcome == CollectionOutcome.FAIL
-    ]
-    if failed_reports:
-        _raise_error_for_failed_generated_tasks(task, failed_reports)
-
-    generated_tasks = [
-        report.node
-        for report in new_reports
-        if report.outcome == CollectionOutcome.SUCCESS
-        and isinstance(report.node, PTask)
-    ]
-    _commit_generated_tasks(session, generated_tasks)
-    # Generators must run to discover tasks even in simulation modes.
-    if session.config["dry_run"] or session.config["explain"]:
-        raise WouldBeExecuted
-    return True
+
+        new_reports: list[CollectionReport] = []
+        for name, function in name_to_function.items():
+            report = session.hook.pytask_collect_task_protocol(
+                session=session,
+                reports=new_reports,
+                path=task.path if isinstance(task, PTaskWithPath) else None,
+                name=name,
+                obj=function,
+            )
+            if report is not None:
+                new_reports.append(report)
+
+        session.collection_reports.extend(new_reports)
+        failed_reports = [
+            report for report in new_reports if report.outcome == CollectionOutcome.FAIL
+        ]
+        if failed_reports:
+            _raise_error_for_failed_generated_tasks(task, failed_reports)
+
+        generated_tasks = [
+            report.node
+            for report in new_reports
+            if report.outcome == CollectionOutcome.SUCCESS
+            and isinstance(report.node, PTask)
+        ]
+        _commit_generated_tasks(session, generated_tasks)
+        # Generators must run to discover tasks even in simulation modes.
+        if session.config["dry_run"] or session.config["explain"]:
+            raise WouldBeExecuted
+        return True
+    finally:
+        _clear_annotation_locals(tasks)
 
 
 def _raise_error_for_failed_generated_tasks(

src/_pytask/task_utils.py

@@ -233,14 +233,18 @@ def create_text_file() -> Annotated[str, Path("file.txt")]:
     ```
 
     """
-    # Capture the caller's frame locals for deferred annotation evaluation in Python
-    # 3.14+. If ``from __future__ import annotations`` is active, keep the pre-3.14
-    # behavior by evaluating annotations against current globals instead of snapshots.
+    # Snapshot definition-local bindings, including stringified annotations. Like
+    # Pydantic, keep future annotations at module scope tied to current globals.
+    # Retain snapshots for native lazy annotations in module-level task loops.
     caller_frame = sys._getframe(1)
     has_future_annotations = bool(
         caller_frame.f_code.co_flags & __future__.annotations.compiler_flag
     )
-    caller_locals = None if has_future_annotations else caller_frame.f_locals.copy()
+    caller_locals = (
+        None
+        if has_future_annotations and caller_frame.f_code.co_name == "<module>"
+        else caller_frame.f_locals.copy()
+    )
 
     def wrapper(func: Callable[P, R_co]) -> TaskDecorated[P, R_co]:
         # Omits frame when a builtin function is wrapped.
@@ -354,6 +358,8 @@ def _parse_after(
         for func in after:
             if not isinstance(func, TaskFunction):
                 func = task()(func)  # noqa: PLW2901
+                # This registration frame is not the callable's definition scope.
+                func.pytask_meta.annotation_locals = None
             new_after.append(func.pytask_meta._id)
         return new_after
     msg = (

tests/conftest.py

@@ -158,6 +158,8 @@ def pytest_collection_modifyitems(session, config, items) -> None:  # noqa: ARG0
         return
     for item in items:
         if isinstance(item, NotebookItem):
+            if item.path.name == "test_task_generator.ipynb":
+                continue
             item.add_marker(pytest.mark.xfail(reason="The tests are flaky."))