@@ -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