ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2227952739
##########
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:
See your earlier review that prompted me to define it here:
https://github.com/apache/airflow/pull/53420#discussion_r2215169496
--
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]
