Lee-W commented on code in PR #43383:
URL: https://github.com/apache/airflow/pull/43383#discussion_r1938884003
##########
airflow/utils/decorators.py:
##########
@@ -18,54 +18,55 @@
from __future__ import annotations
import sys
-from collections import deque
from typing import Callable, TypeVar
+import libcst as cst
+
T = TypeVar("T", bound=Callable)
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name):
+ self.decorators_to_remove = {
Review Comment:
```suggestion
self.decorators_to_remove: set[str] = {
```
##########
providers/standard/tests/provider_tests/standard/utils/test_python_virtualenv.py:
##########
@@ -192,25 +192,25 @@ def
test_should_create_virtualenv_with_extra_packages_uv(self, mock_execute_in_s
)
def test_remove_task_decorator(self):
- py_source = '@task.virtualenv(serializer="dill")\ndef f():\nimport
funcsigs'
+ py_source = '@task.virtualenv(serializer="dill")\ndef f():\n import
funcsigs'
Review Comment:
```suggestion
py_source = '''
@task.virtualenv(serializer="dill")
def f():
import funcsigs
'''
```
we probably should make the test cases look like this way so that we won't
make similar error in the future. (but I guess we have a better way to handle
the indent/dedent thing here 🤔 )
##########
airflow/utils/decorators.py:
##########
@@ -18,54 +18,55 @@
from __future__ import annotations
import sys
-from collections import deque
from typing import Callable, TypeVar
+import libcst as cst
+
T = TypeVar("T", bound=Callable)
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name):
Review Comment:
```suggestion
def __init__(self, task_decorator_name: str) -> None:
```
##########
airflow/utils/decorators.py:
##########
@@ -18,54 +18,55 @@
from __future__ import annotations
import sys
-from collections import deque
from typing import Callable, TypeVar
+import libcst as cst
+
T = TypeVar("T", bound=Callable)
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name):
+ self.decorators_to_remove = {
+ "setup",
+ "teardown",
+ "task.skip_if",
+ "task.run_if",
+ task_decorator_name.strip("@"),
+ }
+
+ def _is_task_decorator(self, decorator: cst.Decorator) -> bool:
+ if isinstance(decorator.decorator, cst.Name):
+ return decorator.decorator.value in self.decorators_to_remove
+ elif isinstance(decorator.decorator, cst.Attribute):
+ if isinstance(decorator.decorator.value, cst.Name):
Review Comment:
```suggestion
elif isinstance(decorator.decorator, cst.Attribute) and
isinstance(decorator.decorator.value, cst.Name):
```
##########
airflow/utils/decorators.py:
##########
@@ -18,54 +18,55 @@
from __future__ import annotations
import sys
-from collections import deque
from typing import Callable, TypeVar
+import libcst as cst
+
T = TypeVar("T", bound=Callable)
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name):
+ self.decorators_to_remove = {
+ "setup",
+ "teardown",
+ "task.skip_if",
+ "task.run_if",
+ task_decorator_name.strip("@"),
+ }
+
+ def _is_task_decorator(self, decorator: cst.Decorator) -> bool:
Review Comment:
```suggestion
def _is_task_decorator(self, decorator_node: cst.Decorator) -> bool:
decorator_expression = decorator_node.decorator
```
I think we can make it more readable here like this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
