sdk: Refactor get_python_source to strip comments (#50327)

[nailo2c]

Closes: #50327

Why

airflow/task-sdk/src/airflow/sdk/definitions/_internal/decorators.py

Lines 76 to 85 in 734960d

	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

The function remove_task_decorator uses cst.parse_module to parse the python source.

However, when the function source contains unusual indentation (e.g., a comment at column-0 inside an indented function), cst.parse_module fails and raises a cst.ParserSyntaxError.

    def b_task():
        print("hello")
##################
        print("more hello")

Because this exception was not caught, the function fails early. This prevents the @task.kubernetes decorator from being removed, causing it to be incorrectly included in the generated /tmp/script.py.

This leads to the final error reported in the issue: NameError: name 'task' is not defined when the pod tries to execute the script.

How

Add a regex fallback for @task.kubernetes and handle the indentation issue.

What

[uranusjr]

I feel this should be fixed in cst instead. Airflow shouldn’t try to work around this on its own.

[amoghrajesh]

I think I agree with @uranusjr here, more of a CST problem that Airflow? Working around would be less than ideal here.

[potiuk]

I think it should be solved by us but elsewhere. The root cause is this function:

    def get_python_source(self):
        raw_source = inspect.getsource(self.python_callable)
        res = textwrap.dedent(raw_source)
        res = remove_task_decorator(res, self.custom_operator_name)
        return res

task-sdk/src/airflow/sdk/bases/decorator.py

We are dedenting the function source ourselves - ans this is where error is introduced.

[potiuk]

And here:

    def get_python_source(self):
        """Return the source of self.python_callable."""
        return textwrap.dedent(inspect.getsource(self.python_callable))

providers/standard/src/airflow/providers/standard/operators/python.py

[potiuk]

We should likely remove the lines containing only comment before dedenting - that should fix the problem

[uranusjr]

Ah that makes sense, nice analysis @potiuk

[amoghrajesh]

Nice investigation @potiuk!

[ashb]

Here's an idea out of left field: don't do any pre-processing or removing at all, but instead put this in the generated script:

from types import SimpleNamespace
task = SimpleNamespace(kubernetes=lambda f: f)

Then we could leave it in the script as:

@task.kubernetes
def my_fn(...): ...

# ...

try:
  my_fn(...)
  ...

That way we might not need to a) use cst to strip the decorator, b) nor worry about dedenting anything?

[nailo2c]

Hi folks, thanks for all the reviews. I believe we now have a better approach version :)

The current version follows @potiuk’s suggestion, it gets rid of lines that start with # in get_python_source.

Please let me know if there's anything else I can improve 💪

[potiuk]

@ashb

from types import SimpleNamespace
task = SimpleNamespace(kubernetes=lambda f: f)

Yeah, my thought exactly that we could not remove the decorators but replace them with something empty. But I think in this case also indentation matter - because we would have to keep the functions at the nested indentation level in the generated script (those callbacks are converted to top-level functions now). I am not sure if that can be handled easily. No esy idea for that one or that it's worth the effort.

I think for now simple removeal of the comments has drawbacks of course. It does not handle all cases - for example it will not work with multi-line strings and likely some other constructs.

    def b_task():
        print("""
hello
"""
        print("more hello")

In order to properly handle those constructs, we would really have to parse the whole code with AST and know which lines are supposed to be indented and which not.

For example this case is not easy to handle without knowing that hello2 is part of a multi-line string -- you need to parse the whole python file to know.

    def b_task():
        print("""
hello

        hello 2
"""
        print("more hello")

Or a better would would be simply to forbid those cases and turn ParserError into a more meaningful one (only use plain functions, do not break indentation. That wouls be simple- we extract the whispace from the first line and reject any function that has line that do not begin with the same whitespace prefix.

That would be simple and effective solution (and back-compatible - those functions now cause ParsingException). We can easily add this limitation and even document it.

[potiuk]

Any other comments @ashb @uranusjr ?