From ecebbdb40378bc9d7502db2489dacfb08b4021bb Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Mon, 12 Oct 2026 09:19:48 +0200 Subject: [PATCH 1/7] Snapshot postponed annotation locals and release them after collection --- .github/workflows/main.yml | 2 +- .github/workflows/zizmor.yml | 2 +- .../how_to_guides/using_task_returns.md | 7 ++ src/_pytask/collect.py | 7 ++ src/_pytask/models.py | 5 +- src/_pytask/task_utils.py | 12 ++- tests/conftest.py | 2 + tests/test_collect.py | 58 +++++++++-- tests/test_jupyter/test_task_generator.ipynb | 4 +- tests/test_provisional.py | 98 ++++++++++++++++++- tests/test_task_utils.py | 25 +++++ 11 files changed, 200 insertions(+), 22 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 41e02aac..d437dac8 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -11,7 +11,7 @@ on: - main pull_request: branches: - - '*' + - '**' merge_group: permissions: {} diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index f394ba6d..8c01f95f 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -6,7 +6,7 @@ on: - main pull_request: branches: - - '*' + - '**' schedule: - cron: '0 7 * * 1' workflow_dispatch: diff --git a/docs/source/how_to_guides/using_task_returns.md b/docs/source/how_to_guides/using_task_returns.md index b6e7ed98..b8eae974 100644 --- a/docs/source/how_to_guides/using_task_returns.md +++ b/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 diff --git a/src/_pytask/collect.py b/src/_pytask/collect.py index 1d4c697d..1f13255a 100644 --- a/src/_pytask/collect.py +++ b/src/_pytask/collect.py @@ -326,6 +326,13 @@ def pytask_collect_task_protocol( else: return None + finally: + # Generated tasks use this protocol without the outer collection hook. + # Release snapshots after both dependencies and products are parsed, even + # when collection fails and the function is retained in an error report. + if isinstance(obj, TaskFunction): + obj.pytask_meta.annotation_locals = None + @hookimpl(trylast=True) def pytask_collect_task( diff --git a/src/_pytask/models.py b/src/_pytask/models.py index 39f31298..ebe63929 100644 --- a/src/_pytask/models.py +++ b/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 diff --git a/src/_pytask/task_utils.py b/src/_pytask/task_utils.py index dd2c1f8c..eb0f5fcb 100644 --- a/src/_pytask/task_utils.py +++ b/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 == "" + 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. diff --git a/tests/conftest.py b/tests/conftest.py index 027f32da..ac684db4 100644 --- a/tests/conftest.py +++ b/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.")) diff --git a/tests/test_collect.py b/tests/test_collect.py index 7a8fcd9b..17a98ce3 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -553,8 +553,13 @@ def task_example(output: Path = Path("in.txt")) -> Annotated[str, output]: assert tmp_path.joinpath("in.txt").read_text() == "Hello" -def test_lazy_annotations_survive_deleted_definition_local(tmp_path): - source = """ +@pytest.mark.parametrize("future_annotations", [False, True]) +def test_lazy_annotations_survive_deleted_definition_local( + tmp_path, future_annotations +): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} from pathlib import Path from typing import Annotated from pytask import task @@ -578,6 +583,33 @@ def task_example() -> Annotated[str, output]: assert tmp_path.joinpath("out.txt").read_text() == "Hello" +@pytest.mark.parametrize("decorator", ["@task", "@task()"]) +def test_nested_future_annotations_use_current_globals(tmp_path, decorator): + source = f""" + from __future__ import annotations + from pathlib import Path + from typing import Annotated + from pytask import task + + OUTPUT = Path("first.txt") + + def create_task(): + {decorator} + def task_example() -> Annotated[str, OUTPUT]: + return "Hello" + + create_task() + OUTPUT = Path("second.txt") + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + + session = build(paths=tmp_path) + + assert session.exit_code == ExitCode.OK + assert tmp_path.joinpath("second.txt").read_text() == "Hello" + assert not tmp_path.joinpath("first.txt").exists() + + def test_string_literal_annotations_are_resolved(tmp_path): source = """ from __future__ import annotations @@ -597,18 +629,26 @@ def task_example() -> 'Annotated[str, OUTPUT]': assert tmp_path.joinpath("out.txt").exists() -def test_annotation_locals_are_cleared_after_collection_to_allow_pickling(tmp_path): - source = """ +@pytest.mark.parametrize("future_annotations", [False, True]) +def test_annotation_locals_are_cleared_after_collection_to_allow_pickling( + tmp_path, future_annotations +): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} import threading from pytask import task - lock = threading.RLock() + def create_tasks(): + lock = threading.RLock() - for i in range(2): - @task - def task_example(): - return None + for i in range(2): + @task + def task_example(): + return None + + create_tasks() """ tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) diff --git a/tests/test_jupyter/test_task_generator.ipynb b/tests/test_jupyter/test_task_generator.ipynb index cb922e56..1407891c 100644 --- a/tests/test_jupyter/test_task_generator.ipynb +++ b/tests/test_jupyter/test_task_generator.ipynb @@ -52,7 +52,9 @@ "outputs": [], "source": [ "session = pytask.build(tasks=[task_create_files, task_generator_copy_files])\n", - "assert session.exit_code == ExitCode.OK" + "assert session.exit_code == ExitCode.OK\n", + "assert Path(\"a-copy.txt\").read_text() == \"Hello, \"\n", + "assert Path(\"b-copy.txt\").read_text() == \"World!\"" ] } ], diff --git a/tests/test_provisional.py b/tests/test_provisional.py index d9140df1..77351a90 100644 --- a/tests/test_provisional.py +++ b/tests/test_provisional.py @@ -2,6 +2,7 @@ import textwrap +import cloudpickle import pytest from _pytask.pluginmanager import get_plugin_manager @@ -165,8 +166,11 @@ def task_depends( assert "There are some tasks which produce" in result.output -def test_provisional_task_generation(runner, tmp_path): - source = """ +@pytest.mark.parametrize("future_annotations", [False, True]) +def test_provisional_task_generation(runner, tmp_path, future_annotations): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} from typing import Annotated from pytask import DirectoryNode, task from pathlib import Path @@ -194,8 +198,94 @@ def task_copy( assert result.exit_code == ExitCode.OK assert "4 Collected tasks" in result.output assert "4 Succeeded" in result.output - assert tmp_path.joinpath("a-copy.txt").exists() - assert tmp_path.joinpath("b-copy.txt").exists() + assert tmp_path.joinpath("a-copy.txt").read_text() == "Hello, " + assert tmp_path.joinpath("b-copy.txt").read_text() == "World!" + + +@pytest.mark.parametrize("future_annotations", [False, True]) +@pytest.mark.parametrize("decorator", ["@task", "@task()"]) +def test_generated_annotations_snapshot_locals_without_parameter_defaults( + tmp_path, future_annotations, decorator +): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} + from pathlib import Path + from typing import Annotated + from pytask import task + + @task(is_generator=True) + def task_generator(): + for i in range(2): + output = Path(f"out-{{i}}.txt") + + {decorator} + def task_example() -> Annotated[str, output]: + return "generated" + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + + session = build(paths=tmp_path) + + assert session.exit_code == ExitCode.OK + assert len(session.tasks) == 3 + assert tmp_path.joinpath("out-0.txt").read_text() == "generated" + assert tmp_path.joinpath("out-1.txt").read_text() == "generated" + + +@pytest.mark.parametrize("fail_collection", [False, True]) +def test_generated_annotation_snapshots_are_released_after_collection( + tmp_path, fail_collection +): + annotation = "unknown" if fail_collection else "output" + source = f""" + from __future__ import annotations + import threading + import weakref + from pathlib import Path + from typing import Annotated + from pytask import task + + class Payload: + pass + + references = [] + + @task(is_generator=True) + def task_generator(): + lock = threading.RLock() + payload = Payload() + references.append(weakref.ref(payload)) + + for i in range(2): + output = Path(f"out-{{i}}.txt") + + @task + def task_example() -> Annotated[str, {annotation}]: + return "generated" + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + + session = build(paths=tmp_path) + + assert session.exit_code == (ExitCode.FAILED if fail_collection else ExitCode.OK) + children = [ + report.node + for report in session.collection_reports + if report.node is not None and "task_example" in report.node.name + ] + assert len(children) == 2 + for child in children: + assert child.function.pytask_meta.annotation_locals is None + cloudpickle.dumps(child.function) + + if not fail_collection: + generator = next( + task for task in session.tasks if task.name.endswith("task_generator") + ) + references = generator.function.__globals__["references"] + assert len(references) == 1 + assert references[0]() is None def test_task_generator_executes_once(tmp_path): diff --git a/tests/test_task_utils.py b/tests/test_task_utils.py index aeb5f2b5..832f4d2b 100644 --- a/tests/test_task_utils.py +++ b/tests/test_task_utils.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path +from typing import Annotated from typing import NamedTuple from uuid import UUID @@ -80,6 +81,30 @@ def task_example(): ... COLLECTED_TASKS.pop(Path(__file__)) +def test_annotation_snapshots_are_shallow_and_preserve_each_loop_binding(): + shared = object() + outputs = [Path("a.txt"), Path("b.txt")] + functions = [] + for output in outputs: + + @task + def task_example() -> Annotated[str, output]: + return "Hello" + + functions.append(task_example) + + first, second = [function.pytask_meta.annotation_locals for function in functions] + assert first is not None + assert second is not None + assert first is not second + assert first["output"] is outputs[0] + assert second["output"] is outputs[1] + assert first["shared"] is shared + assert second["shared"] is shared + + COLLECTED_TASKS.pop(Path(__file__)) + + def test_after_functions_are_stored_as_temporary_ids(): @task def task_before(): ... From 3e21f3279cbc3d081e8ef91c8ce4f488fc82341d Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 09:27:14 +0200 Subject: [PATCH 2/7] Avoid capturing collection locals for programmatic tasks --- src/_pytask/collect.py | 2 ++ tests/test_collect.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/_pytask/collect.py b/src/_pytask/collect.py index 1f13255a..0b31bfe6 100644 --- a/src/_pytask/collect.py +++ b/src/_pytask/collect.py @@ -135,6 +135,8 @@ 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 path = get_file(raw_task) name = raw_task.pytask_meta.name diff --git a/tests/test_collect.py b/tests/test_collect.py index 17a98ce3..3de52b37 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -12,6 +12,7 @@ from _pytask.collect import _find_shortest_uniquely_identifiable_name_for_tasks from _pytask.collect import pytask_collect_node from _pytask.node_protocols import PPathNode +from _pytask.path import import_path from _pytask.pluginmanager import get_plugin_manager from _pytask.task_utils import validate_unique_task_signatures from pytask import CollectionOutcome @@ -532,6 +533,37 @@ def task_example() -> Annotated[str, OUTPUT]: assert not tmp_path.joinpath("first.txt").exists() +def test_programmatic_annotation_globals_do_not_capture_collection_locals(tmp_path): + source = textwrap.dedent( + """ + from __future__ import annotations + from pathlib import Path + from typing import Annotated + + path = Path("out.txt") + + def task_first(): + pass + + def task_second() -> Annotated[str, path]: + return "Hello" + """ + ) + module_path = tmp_path / "functions.py" + module_path.write_text(source) + module = import_path(module_path, tmp_path) + + session = build( + paths=tmp_path, + tasks=[module.task_first, module.task_second], + force=True, + ) + + assert session.exit_code == ExitCode.OK + assert module_path.read_text() == source + assert tmp_path.joinpath("out.txt").read_text() == "Hello" + + def test_annotation_module_globals_take_precedence_over_parameter_defaults(tmp_path): source = """ from __future__ import annotations From 89f968804cb02287a7c1df1b870080afc9e28bf0 Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 09:51:37 +0200 Subject: [PATCH 3/7] Release annotation snapshots after collection hooks finish --- src/_pytask/collect.py | 69 ++++++++++++------------ src/_pytask/provisional.py | 86 ++++++++++++++++-------------- tests/test_collect.py | 106 +++++++++++++++++++++++++++++++++++++ 3 files changed, 188 insertions(+), 73 deletions(-) diff --git a/src/_pytask/collect.py b/src/_pytask/collect.py index 0b31bfe6..d98e6bb8 100644 --- a/src/_pytask/collect.py +++ b/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: @@ -328,13 +338,6 @@ def pytask_collect_task_protocol( else: return None - finally: - # Generated tasks use this protocol without the outer collection hook. - # Release snapshots after both dependencies and products are parsed, even - # when collection fails and the function is retained in an error report. - if isinstance(obj, TaskFunction): - obj.pytask_meta.annotation_locals = None - @hookimpl(trylast=True) def pytask_collect_task( diff --git a/src/_pytask/provisional.py b/src/_pytask/provisional.py index 5b61fb9f..9d78363c 100644 --- a/src/_pytask/provisional.py +++ b/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) @@ -55,47 +57,51 @@ def pytask_execute_task(session: Session, task: PTask) -> bool | None: execute_task(task) - 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( diff --git a/tests/test_collect.py b/tests/test_collect.py index 3de52b37..b0aedfab 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -4,11 +4,13 @@ import textwrap import warnings from pathlib import Path +from typing import Annotated import cloudpickle import pytest import upath +from _pytask._inspect import get_annotations from _pytask.collect import _find_shortest_uniquely_identifiable_name_for_tasks from _pytask.collect import pytask_collect_node from _pytask.node_protocols import PPathNode @@ -585,6 +587,110 @@ def task_example(output: Path = Path("in.txt")) -> Annotated[str, output]: assert tmp_path.joinpath("in.txt").read_text() == "Hello" +@pytest.mark.parametrize("future_annotations", [False, True]) +@pytest.mark.parametrize("is_generator", [False, True]) +def test_annotation_locals_remain_available_to_modify_tasks_hook( + tmp_path, monkeypatch, future_annotations, is_generator +): + future_import = "from __future__ import annotations" if future_annotations else "" + decorator = "@task(is_generator=True)" if is_generator else "" + create_tasks = "" if is_generator else "create_tasks()" + source = f""" + {future_import} + from pathlib import Path + from typing import Annotated + from pytask import task + + {decorator} + def create_tasks(): + output = Path("out.txt") + + @task + def task_example() -> Annotated[str, output]: + return "Hello" + + del output + + {create_tasks} + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + annotations = [] + + class Plugin: + @hookimpl + def pytask_collect_modify_tasks(self, tasks): + annotations.extend( + get_annotations(task.function, eval_str=True)["return"] + for task in tasks + if task.base_name == "task_example" + ) + + pm = get_plugin_manager() + pm.register(Plugin()) + monkeypatch.setattr("_pytask.build.get_plugin_manager", lambda: pm) + + session = build(paths=tmp_path) + + assert session.exit_code == ExitCode.OK + assert annotations == [Annotated[str, Path("out.txt")]] + assert tmp_path.joinpath("out.txt").read_text() == "Hello" + assert all( + task.function.pytask_meta.annotation_locals is None for task in session.tasks + ) + + +@pytest.mark.parametrize("is_generator", [False, True]) +def test_annotation_locals_are_released_when_modify_tasks_hook_fails( + tmp_path, monkeypatch, is_generator +): + decorator = "@task(is_generator=True)" if is_generator else "" + create_tasks = "" if is_generator else "create_tasks()" + source = f""" + from __future__ import annotations + import threading + from pathlib import Path + from typing import Annotated + from pytask import task + + {decorator} + def create_tasks(): + lock = threading.RLock() + output = Path("out.txt") + + @task + def task_example() -> Annotated[str, output]: + return "Hello" + + {create_tasks} + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + functions = [] + + class Plugin: + @hookimpl + def pytask_collect_modify_tasks(self, tasks): + functions.extend( + task.function for task in tasks if task.base_name == "task_example" + ) + if functions: + msg = "Collection hook failed." + raise RuntimeError(msg) + + pm = get_plugin_manager() + pm.register(Plugin()) + monkeypatch.setattr("_pytask.build.get_plugin_manager", lambda: pm) + + session = build(paths=tmp_path) + + assert session.exit_code == ( + ExitCode.FAILED if is_generator else ExitCode.COLLECTION_FAILED + ) + assert len(functions) == 1 + assert functions[0].pytask_meta.annotation_locals is None + cloudpickle.dumps(functions[0]) + assert not tmp_path.joinpath("out.txt").exists() + + @pytest.mark.parametrize("future_annotations", [False, True]) def test_lazy_annotations_survive_deleted_definition_local( tmp_path, future_annotations From 634472203790d6bd9181d4219476b0f458883d67 Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 10:35:38 +0200 Subject: [PATCH 4/7] Document and test programmatic annotation snapshot boundaries --- src/_pytask/collect.py | 2 ++ tests/test_collect.py | 69 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/src/_pytask/collect.py b/src/_pytask/collect.py index d98e6bb8..f8a86285 100644 --- a/src/_pytask/collect.py +++ b/src/_pytask/collect.py @@ -148,6 +148,8 @@ def _collect_from_tasks(session: Session) -> None: # 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 diff --git a/tests/test_collect.py b/tests/test_collect.py index b0aedfab..7b70c974 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -536,6 +536,8 @@ def task_example() -> Annotated[str, OUTPUT]: def test_programmatic_annotation_globals_do_not_capture_collection_locals(tmp_path): + # The fixture's `path` intentionally matches the local in _collect_from_tasks. + # Keep both names in sync to preserve the historical overwrite regression. source = textwrap.dedent( """ from __future__ import annotations @@ -566,6 +568,73 @@ def task_second() -> Annotated[str, path]: assert tmp_path.joinpath("out.txt").read_text() == "Hello" +def test_programmatic_auto_decoration_does_not_capture_locals(tmp_path, monkeypatch): + module_path = tmp_path / "functions.py" + module_path.write_text("def task_example(): pass") + module = import_path(module_path, tmp_path) + snapshots = [] + + class Plugin: + @hookimpl + def pytask_collect_task_setup(self, obj): + snapshots.append(obj.pytask_meta.annotation_locals) + + pm = get_plugin_manager() + pm.register(Plugin()) + monkeypatch.setattr("_pytask.build.get_plugin_manager", lambda: pm) + + session = build(paths=tmp_path, tasks=[module.task_example]) + + assert session.exit_code == ExitCode.OK + assert snapshots == [None] + + +@pytest.mark.parametrize("decorator", ["@task", "@task()"]) +def test_programmatic_collection_preserves_user_annotation_locals( + tmp_path, monkeypatch, decorator +): + source = f""" + from __future__ import annotations + from pathlib import Path + from typing import Annotated + from pytask import task + + def create_task(): + output = Path("out.txt") + + {decorator} + def task_example() -> Annotated[str, output]: + return "Hello" + + return task_example + + task_example = create_task() + """ + module_path = tmp_path / "functions.py" + module_path.write_text(textwrap.dedent(source)) + module = import_path(module_path, tmp_path) + definition_locals = module.task_example.pytask_meta.annotation_locals + assert definition_locals is not None + assert definition_locals["output"] == Path("out.txt") + snapshots = [] + + class Plugin: + @hookimpl + def pytask_collect_task_setup(self, obj): + snapshots.append(obj.pytask_meta.annotation_locals) + + pm = get_plugin_manager() + pm.register(Plugin()) + monkeypatch.setattr("_pytask.build.get_plugin_manager", lambda: pm) + + session = build(paths=tmp_path, tasks=[module.task_example]) + + assert session.exit_code == ExitCode.OK + assert len(snapshots) == 1 + assert snapshots[0] is definition_locals + assert tmp_path.joinpath("out.txt").read_text() == "Hello" + + def test_annotation_module_globals_take_precedence_over_parameter_defaults(tmp_path): source = """ from __future__ import annotations From a8ec4f7a12cc708cdf8dec367b8b3dbceaa50fd9 Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 10:55:36 +0200 Subject: [PATCH 5/7] Avoid capturing internal locals for after prerequisites --- src/_pytask/task_utils.py | 2 ++ tests/test_collect.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/src/_pytask/task_utils.py b/src/_pytask/task_utils.py index eb0f5fcb..83161662 100644 --- a/src/_pytask/task_utils.py +++ b/src/_pytask/task_utils.py @@ -358,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 = ( diff --git a/tests/test_collect.py b/tests/test_collect.py index 7b70c974..93e62cc6 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -535,6 +535,36 @@ def task_example() -> Annotated[str, OUTPUT]: assert not tmp_path.joinpath("first.txt").exists() +@pytest.mark.parametrize("future_annotations", [False, True]) +@pytest.mark.parametrize("annotation_name", ["after", "func", "new_after"]) +def test_after_auto_decoration_does_not_capture_internal_locals( + tmp_path, future_annotations, annotation_name +): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} + from pathlib import Path + from typing import Annotated + from pytask import task + + {annotation_name} = Path("out.txt") + + def task_before() -> Annotated[str, {annotation_name}]: + return "Hello" + + @task(after=task_before) + def task_after(): + pass + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + + session = build(paths=tmp_path) + + assert session.exit_code == ExitCode.OK + assert len(session.tasks) == 2 + assert tmp_path.joinpath("out.txt").read_text() == "Hello" + + def test_programmatic_annotation_globals_do_not_capture_collection_locals(tmp_path): # The fixture's `path` intentionally matches the local in _collect_from_tasks. # Keep both names in sync to preserve the historical overwrite regression. From 3deee3d504a359c6d43326f282808e14eaa8b869 Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 11:12:23 +0200 Subject: [PATCH 6/7] Release child annotation snapshots when generators raise --- src/_pytask/provisional.py | 16 +++++++++++++- tests/test_provisional.py | 44 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/_pytask/provisional.py b/src/_pytask/provisional.py index 9d78363c..f98aea80 100644 --- a/src/_pytask/provisional.py +++ b/src/_pytask/provisional.py @@ -55,7 +55,21 @@ 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 tasks: list[TaskFunction] = [] try: diff --git a/tests/test_provisional.py b/tests/test_provisional.py index 77351a90..5b7c4d08 100644 --- a/tests/test_provisional.py +++ b/tests/test_provisional.py @@ -288,6 +288,50 @@ def task_example() -> Annotated[str, {annotation}]: assert references[0]() is None +@pytest.mark.parametrize("future_annotations", [False, True]) +@pytest.mark.parametrize( + "exception", ["RuntimeError('generator failed')", "SystemExit(7)"] +) +def test_failed_generator_releases_child_annotation_snapshots( + tmp_path, future_annotations, exception +): + future_import = "from __future__ import annotations" if future_annotations else "" + source = f""" + {future_import} + import threading + from pytask import task + + children = [] + error = {exception} + + @task(is_generator=True) + def task_generator(): + lock = threading.RLock() + for i in range(2): + @task + def task_child(): + raise AssertionError("Child must not execute") + children.append(task_child) + raise error + """ + tmp_path.joinpath("task_module.py").write_text(textwrap.dedent(source)) + + session = build(paths=tmp_path) + + assert session.exit_code == ExitCode.FAILED + assert len(session.tasks) == 1 + assert len(session.execution_reports) == 1 + report = session.execution_reports[0] + assert report.outcome == TaskOutcome.FAIL + assert report.exc_info is not None + namespace = report.task.function.__globals__ + assert report.exc_info[1] is namespace["error"] + assert len(namespace["children"]) == 2 + for child in namespace["children"]: + assert child.pytask_meta.annotation_locals is None + cloudpickle.dumps(child) + + def test_task_generator_executes_once(tmp_path): source = """ from pathlib import Path From c79906493c8bb545e20d7f24fe0916103bec21fe Mon Sep 17 00:00:00 2001 From: Tobias Raabe Date: Sat, 17 Oct 2026 11:18:32 +0200 Subject: [PATCH 7/7] docs: add changelog entry for postponed annotation snapshots --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c3d61d55..1782f1ec 100644 --- a/CHANGELOG.md +++ b/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