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/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 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..f8a86285 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: @@ -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 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/provisional.py b/src/_pytask/provisional.py index 5b61fb9f..f98aea80 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) @@ -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( diff --git a/src/_pytask/task_utils.py b/src/_pytask/task_utils.py index dd2c1f8c..83161662 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. @@ -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 = ( 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..93e62cc6 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -4,14 +4,17 @@ 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 +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 +535,136 @@ 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. + 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_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 @@ -553,8 +686,117 @@ 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]) +@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 +): + 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 +820,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 +866,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..5b7c4d08 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,138 @@ 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 + + +@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): 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(): ...