amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2238625812
##########
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:
Actually, this test isn't doing anything too much, let's move this test into
`task-sdk/tests/task_sdk/definitions/decorators/`
--
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]
