danvics/ytdl-sub
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

better union output types

Browse source
This commit is contained in:
Jesse Bannon 2023-12-06 23:50:17 -08:00
parent 9004fba5cb
commit 2ffc53fc07
2 changed files with 40 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
src/ytdl_sub/script/types/function.py
View file

@ -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,

23
tests/unit/script/test_parser.py
View file

@ -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(