branching + dyn. task mapping + passing arguments: how to filter passed arguments

[olk]

Unfortunately all branches get the same data/arguments passed.
Based on the branch task-group 'consumer_a' should only process parameters with "flag" = True and 'consumer_b' only data with "flag" = False.
How can I filter the argument list?

from airflow.decorators import dag, task, task_group
from pendulum import datetime


@task
def producer():
    return [ { "flag": True, "url": "abc" }, { "flag": False, "url": "xyz "} ]


@task.branch
def branching(data):
    if data["flag"]:
        return "consumer_a"
    else:
        return "consumer_b"


@task
def print_flag(data):
    print(f"flag: {data['flag']}")
    return data


@task
def print_url(data):
    print(f"url: {data['url']}")


@task_group
def consumer_a(data):
    print_url(print_flag(data))


@task_group
def consumer_b(data):
    print_url(print_flag(data))


@dag(
    start_date=datetime(2023, 1, 1),
    catchup=False,
    schedule="@daily"
)
def branch_2():

    data = producer()
    random_choice.expand(data=data) >> [ consumer_a.expand(data=data), consumer_b.expand(data=data) ] # filtering 'data' here?


branch_2()

[olk]

A variant that solves the problem doe expand the arguments before branching but I'm wondering if there is a more elegant solution ...

from airflow.decorators import dag, task, task_group
from pendulum import datetime


@task
def producer():
    return [ { "flag": True, "url": "abc" }, { "flag": False, "url": "xyz "} ]


@task.branch
def branching(data):
    if data["flag"]:
        return "expanding.consumer_a"
    else:
        return "expanding.consumer_b"


@task
def print_flag(data):
    print(f"flag: {data['flag']}")
    return data


@task
def print_url(data):
    print(f"url: {data['url']}")


@task_group
def consumer_a(data):
    print_url(print_flag(data))


@task_group
def consumer_b(data):
    print_url(print_flag(data))


@task_group
def expanding(data):
    branching(data=data) >> [ consumer_a(data=data), consumer_b(data=data) ]


@dag(
    start_date=datetime(2023, 1, 1),
    catchup=False,
    schedule="@daily"
)
def branch_2():

    data = producer()
    expanding.expand(data=data)


branch_2()