amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2215156974
##########
airflow-core/src/airflow/models/deadline.py:
##########
@@ -44,6 +43,33 @@
logger = logging.getLogger(__name__)
+class classproperty:
+ """
+ Decorator that converts a method with a single cls argument into a
property.
+
+ Mypy won't let us use both @property and @classmethod together, this is a
workaround
+ to combine the two.
+
+ Usage:
+
+ class Circle:
+ def __init__(self, radius):
+ self.radius = radius
+
+ @classproperty
+ def pi(cls):
+ return 3.14159
+
+ print(Circle.pi) # Outputs: 3.14159
+ """
+
+ def __init__(self, method):
+ self.method = method
+
+ def __get__(self, instance, cls=None):
+ return self.method(cls)
Review Comment:
Cool
##########
airflow-core/tests/unit/utils/test_preexisting_python_virtualenv_decorator.py:
##########
@@ -21,7 +21,7 @@
import pytest
-from airflow.utils.decorators import remove_task_decorator
+from airflow.sdk.bases.decorator import remove_task_decorator
Review Comment:
We shouldn't be importing from sdk here, created a dependency which we are
trying to chop.
This function could probably be easily duplicated.
##########
task-sdk/src/airflow/sdk/bases/decorator.py:
##########
@@ -684,3 +685,66 @@ def decorator_factory(python_callable):
)
return cast("TaskDecorator", decorator_factory)
+
+
+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
+
+
+class _autostacklevel_warn:
+ def __init__(self, delta):
+ self.warnings = __import__("warnings")
+ self.delta = delta
+
+ def __getattr__(self, name):
+ return getattr(self.warnings, name)
+
+ def __dir__(self):
+ return dir(self.warnings)
+
+ def warn(self, message, category=None, stacklevel=1, source=None):
+ self.warnings.warn(message, category, stacklevel + self.delta, source)
+
+
+def fixup_decorator_warning_stack(func, delta: int = 2):
+ if func.__globals__.get("warnings") is sys.modules["warnings"]:
+ # Yes, this is more than slightly hacky, but it _automatically_ sets
the right stacklevel parameter to
+ # `warnings.warn` to ignore the decorator.
+ func.__globals__["warnings"] = _autostacklevel_warn(delta)
Review Comment:
We already have this defined under:
`task-sdk/src/airflow/sdk/definitions/_internal/decorators.py`
##########
providers/standard/tests/unit/standard/utils/test_python_virtualenv.py:
##########
@@ -24,7 +24,12 @@
import pytest
from airflow.providers.standard.utils.python_virtualenv import
_generate_pip_conf, _use_uv, prepare_virtualenv
-from airflow.utils.decorators import remove_task_decorator
+
+try:
+ from airflow.sdk.bases.decorator import remove_task_decorator
+except ImportError:
+ # Fallback for Airflow < 3.1
+ from airflow.utils.decorators import remove_task_decorator # type:
ignore[no-redef]
Review Comment:
Since this is a test, the version_compat under
`devel-common/src/tests_common/test_utils/version_compat.py`
##########
airflow-core/src/airflow/utils/__init__.py:
##########
@@ -46,5 +46,9 @@ def __getattr__(name: str):
"xcom": {
"XCOM_RETURN_KEY": "airflow.models.xcom.XCOM_RETURN_KEY",
},
+ "decorators": {
+ "remove_task_decorator":
"airflow.sdk.bases.decorator.remove_task_decorator",
Review Comment:
Since all of these are internal, you should move them to:
`task-sdk/src/airflow/sdk/definitions/_internal/decorators.py` imo.
##########
providers/standard/tests/unit/standard/utils/test_python_virtualenv.py:
##########
@@ -24,7 +24,12 @@
import pytest
from airflow.providers.standard.utils.python_virtualenv import
_generate_pip_conf, _use_uv, prepare_virtualenv
-from airflow.utils.decorators import remove_task_decorator
+
+try:
+ from airflow.sdk.bases.decorator import remove_task_decorator
+except ImportError:
+ # Fallback for Airflow < 3.1
+ from airflow.utils.decorators import remove_task_decorator # type:
ignore[no-redef]
Review Comment:
Could you add this to the version_compat module and import it?
--
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]
