You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ephraimbuddy (via GitHub)" <gi...@apache.org> on 2025/07/16 13:49:16 UTC
[PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
ephraimbuddy opened a new pull request, #53420:
URL: https://github.com/apache/airflow/pull/53420
`classproperty` was only used in airflow.models.deadline.py so I moved it there. We can find a better place for it when there are more use of it outside the deadline module
Closes: #53419
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2222596313
##########
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:
It's already backwards compatible by virtue of adding it in this place. A test:
<img width="842" height="103" alt="Screenshot 2025-07-22 at 14 43 53" src="https://github.com/user-attachments/assets/dd18f073-b5d5-4c04-98fa-1763672cba4b" />
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2215766519
##########
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:
Sorry, I mean to say public not internal
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2228049725
##########
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:
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2215218561
##########
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:
I think the remove_task_decorator is technically internal or am I missing something?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2218422124
##########
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:
It doesn't look like it has use aside tests, however, no one knows what people can use in the wild but let's make it private
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2230388372
##########
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:
It was repeated since you suggested it in an earlier review. Should we do something different here?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2221344003
##########
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:
Agreed, maybe add a backcompat just to be sure or it becomes a breaking change?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy merged PR #53420:
URL: https://github.com/apache/airflow/pull/53420
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2218334402
##########
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:
Does it have a need to have it public? If there is no need to have it public, let's make it internal
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #53420:
URL: https://github.com/apache/airflow/pull/53420#discussion_r2215226477
##########
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:
It is internal yeah, that's why I am suggesting we move it into: `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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Move functions in `airflow.utils.decorator` to more appropriate places [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org