Improve stability of remove_task_decorator function (#38649)

* Improve stability of remove_task_decorator function * fix statics * test * remove test --------- Co-authored-by: Sam Wheating <[email protected]>
apache · Apr 1, 2024 · f1301da · f1301da
1 parent d3dc88f
commit f1301da
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/airflow/utils/decorators.py b/airflow/utils/decorators.py
@@ -62,12 +62,18 @@ def remove_task_decorator(python_source: str, task_decorator_name: str) -> str:
 
     :param python_source: The python source code
     :param task_decorator_name: the decorator name
+
+    TODO: Python 3.9+: Rewrite this to use ast.parse and ast.unparse
     """
 
     def _remove_task_decorator(py_source, decorator_name):
-        if decorator_name not in py_source:
+        # if no line starts with @decorator_name, we can early exit
+        for line in py_source.split("\n"):
+            if line.startswith(decorator_name):
+                break
+        else:
             return python_source
-        split = python_source.split(decorator_name)
+        split = python_source.split(decorator_name, 1)
         before_decorator, after_decorator = split[0], split[1]
         if after_decorator[0] == "(":
             after_decorator = _balance_parens(after_decorator)

diff --git a/tests/utils/test_python_virtualenv.py b/tests/utils/test_python_virtualenv.py
@@ -125,6 +125,11 @@ def test_remove_decorator_no_parens(self):
         res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv")
         assert res == "def f():\nimport funcsigs"
 
+    def test_remove_decorator_including_comment(self):
+        py_source = "@task.virtualenv\ndef f():\n# @task.virtualenv\nimport funcsigs"
+        res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv")
+        assert res == "def f():\n# @task.virtualenv\nimport funcsigs"
+
     def test_remove_decorator_nested(self):
         py_source = "@foo\n@task.virtualenv\n@bar\ndef f():\nimport funcsigs"
         res = remove_task_decorator(python_source=py_source, task_decorator_name="@task.virtualenv")