From 2a2f9780744dd38f98fa200f4c157471fec572a9 Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Thu, 5 Dec 2024 17:34:52 +0000
Subject: [PATCH] Use + as prefix for taskflow references

---
 dagfactory/dagbuilder.py      | 6 +++---
 dev/dags/example_taskflow.yml | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py
index 58675d54..4dd1ff2c 100644
--- a/dagfactory/dagbuilder.py
+++ b/dagfactory/dagbuilder.py
@@ -987,15 +987,15 @@ def make_decorator(
                 else operator_obj.partial(**task_params).expand(**expand_kwargs)
             )
         """
-
         callable_args_keys = inspect.getfullargspec(python_callable).args
         callable_kwargs = {}
         decorator_kwargs = dict(**task_params)
         for arg_key, arg_value in task_params.items():
             if arg_key in callable_args_keys:
                 decorator_kwargs.pop(arg_key)
-                if arg_value in tasks_dict:
-                    callable_kwargs[arg_key] = tasks_dict[arg_value]
+                if arg_value.startswith("+"):
+                    upstream_task_name = arg_value.split("+")[0]
+                    callable_kwargs[arg_key] = tasks_dict[upstream_task_name]
                 else:
                     callable_kwargs[arg_key] = arg_value
 
diff --git a/dev/dags/example_taskflow.yml b/dev/dags/example_taskflow.yml
index c1e407d8..60c8cb91 100644
--- a/dev/dags/example_taskflow.yml
+++ b/dev/dags/example_taskflow.yml
@@ -8,19 +8,19 @@ example_taskflow:
   tasks:
     some_number:
       decorator: airflow.decorators.task
-      python_callable: "sample.some_number"
+      python_callable: sample.some_number
     numbers_list:
       decorator: airflow.decorators.task
       python_callable_name: build_numbers_list
       python_callable_file: $CONFIG_ROOT_DIR/sample.py
     double_number_from_arg:
       decorator: airflow.decorators.task
-      python_callable: "sample.double"
+      python_callable: sample.double
       number: 2
     double_number_from_task:
       decorator: airflow.decorators.task
-      python_callable: "sample.double"
-      number: some_number  # this is a task previously defined
+      python_callable: sample.double
+      number: +some_number  # this is a task previously defined
       #dependencies: [some_number]
     #double_list_dynamically:
     #  python_callable: !!python/name:sample.double"