better union output types

src/ytdl_sub/script/types/function.py

@@ -107,22 +107,26 @@ class BuiltInFunction(Function, BuiltInFunctionType):
             return arg.output_type()
         return type(arg)
 
-    def output_type(self) -> Type[Resolvable]:
+    def _output_generic_type(self, union_args: List[Type[Argument]]) -> Type[Resolvable]:
-        if self.function_spec.return_type is ReturnableArgument:
+        union_types_list = set()
-            generic_arg_index = self.function_spec.args.index(ReturnableArgument)
+        for union_type in union_args:
-            return self._arg_output_type(self.args[generic_arg_index])
+            possible_output_type = union_type
-        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]))
+                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.append(union_type)
+                union_types_list.add(possible_output_type)
 
         return Union[tuple(union_types_list)]
 
-        return self.function_spec.return_type
+    def output_type(self) -> Type[Resolvable]:
+        if is_union(self.function_spec.return_type):
+            return self._output_generic_type(self.function_spec.return_type.__args__)
+
+        return self._output_generic_type([self.function_spec.return_type])
 
     def _resolve_lambda_function(
         self,

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(