Fix tuple handling in condition branches

flytekit/annotated/promise.py

@@ -384,6 +384,13 @@ def with_overrides(self, *args, **kwargs):
  val.with_overrides(*args, **kwargs)
  return self
 
+ @property
+ def ref(self):
+ for p in promises:
+ if p.ref:
+ return p.ref
+ return None
+
  return Output(*promises)
 
 

tests/flytekit/unit/annotated/test_conditions.py

@@ -1,3 +1,4 @@
+import typing
 import pytest
 
 from flytekit import task, workflow
@@ -47,3 +48,57 @@ def multiplier_2(my_input: float) -> float:
 
  with pytest.raises(ValueError):
  multiplier_2(my_input=10)
+
+
+def test_condition_sub_workflows():
+ @task()
+ def sum_div_sub(a: int, b: int) -> typing.NamedTuple("Outputs", sum=int, div=int, sub=int):
+ return a + b, a / b, a - b
+
+ @task()
+ def sum_sub(a: int, b: int) -> typing.NamedTuple("Outputs", sum=int, sub=int):
+ return a + b, a - b
+
+ @workflow
+ def sub_wf(a: int, b: int) -> (int, int):
+ return sum_sub(a=a, b=b)
+
+ @workflow
+ def math_ops(a: int, b: int) -> (int, int):
+ # Flyte will only make `sum` and `sub` available as outputs because they are common between all branches
+ sum, sub = (
+ conditional("noDivByZero")
+ .if_(a > b)
+ .then(sub_wf(a=a, b=b))
+ .else_()
+ .fail("Only positive results are allowed")
+ )
+
+ return sum, sub
+
+ x, y = math_ops(a=3, b=2)
+ assert x == 5
+ assert y == 1
+
+
+def test_condition_tuple_branches():
+ @task()
+ def sum_sub(a: int, b: int) -> typing.NamedTuple("Outputs", sum=int, sub=int):
+ return a + b, a - b
+
+ @workflow
+ def math_ops(a: int, b: int) -> (int, int):
+ # Flyte will only make `sum` and `sub` available as outputs because they are common between all branches
+ sum, sub = (
+ conditional("noDivByZero")
+ .if_(a > b)
+ .then(sum_sub(a=a, b=b))
+ .else_()
+ .fail("Only positive results are allowed")
+ )
+
+ return sum, sub
+
+ x, y = math_ops(a=3, b=2)
+ assert x == 5
+ assert y == 1