From 2ffc53fc0729fe6e63efd3e5e2e861e97b88fd59 Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Wed, 6 Dec 2023 23:50:17 -0800
Subject: [PATCH] better union output types

---
 src/ytdl_sub/script/types/function.py | 30 +++++++++++++++------------
 tests/unit/script/test_parser.py      | 23 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py
index 40fa95c8..bc50d5ad 100644
--- a/src/ytdl_sub/script/types/function.py
+++ b/src/ytdl_sub/script/types/function.py
@@ -107,22 +107,26 @@ class BuiltInFunction(Function, BuiltInFunctionType):
             return arg.output_type()
         return type(arg)
 
+    def _output_generic_type(self, union_args: List[Type[Argument]]) -> Type[Resolvable]:
+        union_types_list = set()
+        for union_type in union_args:
+            possible_output_type = union_type
+            if union_type in (ReturnableArgument, ReturnableArgumentA, ReturnableArgumentB):
+                generic_arg_index = self.function_spec.args.index(union_type)
+                possible_output_type = self._arg_output_type(self.args[generic_arg_index])
+
+            if is_union(possible_output_type):
+                union_types_list.update(possible_output_type.__args__)
+            else:
+                union_types_list.add(possible_output_type)
+
+        return Union[tuple(union_types_list)]
+
     def output_type(self) -> Type[Resolvable]:
-        if self.function_spec.return_type is ReturnableArgument:
-            generic_arg_index = self.function_spec.args.index(ReturnableArgument)
-            return self._arg_output_type(self.args[generic_arg_index])
         if is_union(self.function_spec.return_type):
-            union_types_list = []
-            for union_type in self.function_spec.return_type.__args__:
-                if union_type in (ReturnableArgument, ReturnableArgumentA, ReturnableArgumentB):
-                    generic_arg_index = self.function_spec.args.index(union_type)
-                    union_types_list.append(self._arg_output_type(self.args[generic_arg_index]))
-                else:
-                    union_types_list.append(union_type)
+            return self._output_generic_type(self.function_spec.return_type.__args__)
 
-            return Union[tuple(union_types_list)]
-
-        return self.function_spec.return_type
+        return self._output_generic_type([self.function_spec.return_type])
 
     def _resolve_lambda_function(
         self,
diff --git a/tests/unit/script/test_parser.py b/tests/unit/script/test_parser.py
index 6a922636..196ae8f2 100644
--- a/tests/unit/script/test_parser.py
+++ b/tests/unit/script/test_parser.py
@@ -78,6 +78,29 @@ class TestParser:
             ]
         )
 
+    def test_nested_if_output_type(self):
+        parsed = parse(
+            """{
+            %if(
+                True,
+                %if(
+                    True,
+                    %if(
+                        True,
+                        "winner",
+                        True
+                    ),
+                    True
+                ),
+                True
+            )
+        }"""
+        )
+        assert len(parsed.ast) == 1
+        token = parsed.ast[0]
+        assert isinstance(token, BuiltInFunction)
+        assert token.output_type() == Union[String, Boolean]
+
     def test_single_function_one_vararg(self):
         parsed = parse("hello {%concat('hi mom')}")
         assert parsed == SyntaxTree(