Allow use of callable for `template_fields` for TaskFlow API tasks

[matthewblock]

Description

PR #37028 in Airflow 2.10.0 introduced the ability to pass a Python callable to templateable fields, instead of a Jinja templated string:

https://www.astronomer.io/docs/learn/templating/#use-a-python-callable-for-template-fields

In Airflow 2.10+ it is possible to pass a Python callable to templateable fields. This is especially useful when the parameter value is created using complex operations that might not be possible or are hard to read in Jinja.

However, this feature doesn't seem to work for @task.* decorated tasks.

IMO it doesn't make sense to introduce a new feature that only works with the old style of tasks.

This feature is to make the new Python callable templating work for @task decorated TaskFlow API tasks.

Use case/motivation

I want to specify a Python function for a templated field, instead of a Jinja templated string, for the reasons mentioned in the linked PR. For complex logic, it's much easier to implement and read Jinja rendering functions than a Jinja templated string.

Related issues

#35844
#37028

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[topherinternational]

Could you provide a code snippet of a @task-decorated function where the behavior isn't working as you expect?

[potiuk]

Task flow operators do not have templateable field at all - so it does not make sense to add it there. And in task flow task you can simply call a Python function from inside the decorated function - there is no need to pass callable to the taskflow parameter if you can call them directly. This request makes no sense.

[matthewblock]

Task flow operators do not have templateable field at all

Maybe my terminology is off, but if I decorate PythonOperator, the following are still templateable fields, because the user can specify a Jinja template string and it will get rendered by the time the task executes.

• op_kwargs
• op_args
• templates_dict

The first two technically can't be specified in a decorated task, instead Airflow is rolling up any "extra" args or kwargs the user specifies when they instantiate a task, but they still get rendered.

Here is a test DAG that demonstrates what I'm talking about - See the last task, which prints the un-rendered function.

from datetime import datetime, timedelta

from airflow.decorators import dag, task

default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "retry_delay": timedelta(minutes=1),
    "retries": 0,
}

@dag(
    default_args=default_args,
    start_date=datetime(2000, 1, 1),
    catchup=False,
    schedule=None,
)
def test_jinja_functions_dag():
    def get_foo(context, jinja_env):
        return "bar"

    @task
    def args_kwargs_task(*args, **kwargs):
        print(args) # Prints tuple containing current date (ds)
        print(kwargs["foo"]) # Prints Task Instance object (ti)

    @task(templates_dict={"some_thing": "{{ ds }}"})
    def templates_dict_class(**kwargs):
        print(kwargs["templates_dict"]["some_thing"]) # Prints current date (ds)

    @task
    def direct_kwarg_task(foo):
        print(foo) # Prints Task Instance object (ti)

    @task
    def direct_kwarg_no_rendering_task(foo):
        print(foo) # Prints function object get_foo, NOT "bar"

    (
        args_kwargs_task("{{ ds }}", foo="{{ ti }}"),
        templates_dict_class(),
        direct_kwarg_task(foo="{{ ti }}"),
        direct_kwarg_no_rendering_task(foo=get_foo),
    )

test_jinja_functions_dag()

And in task flow task you can simply call a Python function from inside the decorated function - there is no need to pass callable to the taskflow parameter if you can call them directly.

You can also do this in your python_callable in a non-decorated PythonOperator task 🤷‍♂️ but that's not the point. As a user, it doesn't make sense why I'm encouraged to use TaskFlow API going forward, but I can only use a new syntax introduced in 2.10.0 using the old style PythonOperator.

When I was developing my DAG, I was surprised that my template function didn't actually get called; my task didn't fail because my function outputted a string and instead the name of the function <function test_jinja_functions_dag.<***s>.get_foo at 0xffffacac0a40> was getting passed everywhere.

At the least I think the documentation should be updated.

See references:

• https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/operators.html#jinja-templating
• https://airflow.apache.org/docs/apache-airflow/stable/tutorial/taskflow.html#using-templates-in-decorated-tasks
• https://stackoverflow.com/questions/70134140/how-can-i-pass-jinja2-templated-parameter-to-a-task-using-the-taskflow-api
• https://stackoverflow.com/questions/79055995/airflow-python-operator-typeerror-got-multiple-values-for-keyword-argument-op

[matthewblock]

@potiuk can I open this as a discussion? Maybe this doesn't warrant a feature request but IMO this behavior is unexpected