amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2227475641
##########
task-sdk/src/airflow/sdk/definitions/_internal/decorators.py:
##########
@@ -18,25 +18,68 @@
from __future__ import annotations
import sys
-from types import FunctionType
+
+import libcst as cst
class _autostacklevel_warn:
- def __init__(self):
+ def __init__(self, delta):
self.warnings = __import__("warnings")
+ self.delta = delta
- def __getattr__(self, name: str):
+ def __getattr__(self, name):
Review Comment:
Lets keep the type hint?
##########
airflow-core/tests/unit/utils/test_preexisting_python_virtualenv_decorator.py:
##########
@@ -19,9 +19,48 @@
from textwrap import dedent
+import libcst as cst
import pytest
-from airflow.utils.decorators import remove_task_decorator
+
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name: str) -> None:
+ self.decorators_to_remove: set[str] = {
+ "setup",
+ "teardown",
+ "task.skip_if",
+ "task.run_if",
+ task_decorator_name.strip("@"),
+ }
+
+ def _is_task_decorator(self, decorator_node: cst.Decorator) -> bool:
+ decorator_expr = decorator_node.decorator
+ if isinstance(decorator_expr, cst.Name):
+ return decorator_expr.value in self.decorators_to_remove
+ if isinstance(decorator_expr, cst.Attribute) and
isinstance(decorator_expr.value, cst.Name):
+ return f"{decorator_expr.value.value}.{decorator_expr.attr.value}"
in self.decorators_to_remove
+ if isinstance(decorator_expr, cst.Call):
+ return
self._is_task_decorator(cst.Decorator(decorator=decorator_expr.func))
+ return False
+
+ def leave_FunctionDef(
+ self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+ ) -> cst.FunctionDef:
+ new_decorators = [dec for dec in updated_node.decorators if not
self._is_task_decorator(dec)]
+ if len(new_decorators) == len(updated_node.decorators):
+ return updated_node
+ return updated_node.with_changes(decorators=new_decorators)
Review Comment:
In fact this whole thing is repeated?
##########
airflow-core/tests/unit/utils/test_preexisting_python_virtualenv_decorator.py:
##########
@@ -19,9 +19,48 @@
from textwrap import dedent
+import libcst as cst
import pytest
-from airflow.utils.decorators import remove_task_decorator
+
+class _TaskDecoratorRemover(cst.CSTTransformer):
+ def __init__(self, task_decorator_name: str) -> None:
+ self.decorators_to_remove: set[str] = {
+ "setup",
+ "teardown",
+ "task.skip_if",
+ "task.run_if",
+ task_decorator_name.strip("@"),
+ }
+
+ def _is_task_decorator(self, decorator_node: cst.Decorator) -> bool:
+ decorator_expr = decorator_node.decorator
+ if isinstance(decorator_expr, cst.Name):
+ return decorator_expr.value in self.decorators_to_remove
+ if isinstance(decorator_expr, cst.Attribute) and
isinstance(decorator_expr.value, cst.Name):
+ return f"{decorator_expr.value.value}.{decorator_expr.attr.value}"
in self.decorators_to_remove
+ if isinstance(decorator_expr, cst.Call):
+ return
self._is_task_decorator(cst.Decorator(decorator=decorator_expr.func))
+ return False
+
+ def leave_FunctionDef(
+ self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+ ) -> cst.FunctionDef:
+ new_decorators = [dec for dec in updated_node.decorators if not
self._is_task_decorator(dec)]
+ if len(new_decorators) == len(updated_node.decorators):
+ return updated_node
+ return updated_node.with_changes(decorators=new_decorators)
+
+
+def remove_task_decorator(python_source: str, task_decorator_name: str) -> str:
+ """
+ Remove @task or similar decorators as well as @setup and @teardown.
+ :param python_source: The python source code
+ :param task_decorator_name: the decorator name
+ """
+ source_tree = cst.parse_module(python_source)
+ modified_tree =
source_tree.visit(_TaskDecoratorRemover(task_decorator_name))
+ return modified_tree.code
Review Comment:
It's interesting to me why this is redefined here, we already have it in
`task-sdk/src/airflow/sdk/definitions/_internal/decorators.py`
--
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]
