Skip to content

Commit 192856b

Browse files
committed
Remove _caller_locals parameter from task decorator
Instead of recursively calling task() and passing captured locals via a parameter, detect the no-parentheses case early and call wrapper() directly. The closure captures caller_locals, eliminating the need for the parameter. This simplifies the public API while maintaining the same behavior for deferred annotation evaluation in Python 3.14+.
1 parent 9231432 commit 192856b

1 file changed

Lines changed: 20 additions & 16 deletions

File tree

src/_pytask/task_utils.py

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@ def task( # noqa: PLR0913
5252
id: str | None = None, # noqa: A002
5353
kwargs: dict[Any, Any] | None = None,
5454
produces: Any | None = None,
55-
_caller_locals: dict[str, Any] | None = None,
5655
) -> Callable[..., Callable[..., Any]]:
5756
"""Decorate a task function.
5857
@@ -99,22 +98,29 @@ def create_text_file() -> Annotated[str, Path("file.txt")]:
9998
10099
"""
101100
# Capture the caller's frame locals for deferred annotation evaluation in Python
102-
# 3.14+. This must be done here (not in wrapper) to get the correct scope when
103-
# @task is used without parentheses.
104-
if _caller_locals is None:
105-
_caller_locals = sys._getframe(1).f_locals.copy()
106-
107-
def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
108-
# Omits frame when a builtin function is wrapped.
109-
_rich_traceback_omit = True
101+
# 3.14+. The wrapper closure captures this variable.
102+
caller_locals = sys._getframe(1).f_locals.copy()
110103

104+
# Detect if decorator is used without parentheses: @task instead of @task()
105+
# In this case, `name` is actually the function being decorated.
106+
if is_task_function(name) and kwargs is None:
107+
func_to_wrap = name
108+
actual_name = None
109+
else:
110+
func_to_wrap = None
111+
actual_name = name
112+
# Validate arguments only when used with parentheses
111113
for arg, arg_name in ((name, "name"), (id, "id")):
112114
if not (isinstance(arg, str) or arg is None):
113115
msg = (
114116
f"Argument {arg_name!r} of @task must be a str, but it is {arg!r}."
115117
)
116118
raise ValueError(msg)
117119

120+
def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
121+
# Omits frame when a builtin function is wrapped.
122+
_rich_traceback_omit = True
123+
118124
unwrapped = unwrap_task_function(func)
119125
if isinstance(unwrapped, Function):
120126
coiled_kwargs = extract_coiled_function_kwargs(unwrapped)
@@ -135,7 +141,7 @@ def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
135141
path = get_file(unwrapped)
136142

137143
parsed_kwargs = {} if kwargs is None else kwargs
138-
parsed_name = _parse_name(unwrapped, name)
144+
parsed_name = _parse_name(unwrapped, actual_name)
139145
parsed_after = _parse_after(after)
140146

141147
if hasattr(unwrapped, "pytask_meta"):
@@ -146,11 +152,11 @@ def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
146152
unwrapped.pytask_meta.markers.append(Mark("task", (), {}))
147153
unwrapped.pytask_meta.name = parsed_name
148154
unwrapped.pytask_meta.produces = produces
149-
unwrapped.pytask_meta.annotation_locals = _caller_locals
155+
unwrapped.pytask_meta.annotation_locals = caller_locals
150156
else:
151157
unwrapped.pytask_meta = CollectionMetadata( # type: ignore[attr-defined]
152158
after=parsed_after,
153-
annotation_locals=_caller_locals,
159+
annotation_locals=caller_locals,
154160
is_generator=is_generator,
155161
id_=id,
156162
kwargs=parsed_kwargs,
@@ -168,10 +174,8 @@ def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
168174

169175
return unwrapped
170176

171-
# In case the decorator is used without parentheses, wrap the function which is
172-
# passed as the first argument with the default arguments.
173-
if is_task_function(name) and kwargs is None:
174-
return task(_caller_locals=_caller_locals)(name)
177+
if func_to_wrap is not None:
178+
return wrapper(func_to_wrap)
175179
return wrapper
176180

177181

0 commit comments

Comments
 (0)