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

[BUGFIX] Do not run all branches of if functions in scripts (#999)

Browse source 
Will only run branch script code of `if` statements if the condition evaluates to that branch
This commit is contained in:
Jesse Bannon 2024-06-03 15:02:59 -07:00 committed by GitHub
parent 5e335b195c
commit 5866c12104
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 195 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
src/ytdl_sub/script/functions/conditional_functions.py
View file

@ -18,8 +18,8 @@ class ConditionalFunctions:
depending on the ``condition`` value.
"""
if condition.value:
return true
return false
return true.value()
return false.value()
@staticmethod
def elif_(*if_elif_else: AnyArgument) -> AnyArgument:
@ -50,9 +50,9 @@ class ConditionalFunctions:
for idx in range(0, len(arguments) - 1, 2):
if bool(arguments[idx].value):
return arguments[idx + 1]
return arguments[idx + 1].value()
return arguments[-1]
return arguments[-1].value()
@staticmethod
def if_passthrough(
@ -63,6 +63,7 @@ class ConditionalFunctions:
Conditional ``if`` statement that returns the ``maybe_true_arg`` if it evaluates to True,
otherwise returns ``else_arg``.
"""
if bool(maybe_true_arg.value):
return maybe_true_arg
return else_arg
maybe_true_value = maybe_true_arg.value()
if bool(maybe_true_value.value):
return maybe_true_value
return else_arg.value()

42
src/ytdl_sub/script/script.py
View file

@ -7,6 +7,7 @@ from typing import Set
from ytdl_sub.script.functions import Functions
from ytdl_sub.script.parser import parse
from ytdl_sub.script.script_output import ScriptOutput
from ytdl_sub.script.types.resolvable import BuiltInFunctionType
from ytdl_sub.script.types.resolvable import Lambda
from ytdl_sub.script.types.resolvable import Resolvable
from ytdl_sub.script.types.syntax_tree import SyntaxTree
@ -142,22 +143,47 @@ class Script:
f"{nested_custom_function.num_input_args}"
)
def _get_lambda_function_names_to_evaluate(self, function: BuiltInFunctionType) -> Set[str]:
lambda_function_names: Set[str] = set()
for lamb in SyntaxTree(function.args).lambdas:
if lamb in function.args:
lambda_function_names.add(lamb.value)
# See if the arg outputs a lambda (from an if).
# If so, add the possible lambda to be checked
for arg in function.args:
if (
isinstance(arg, BuiltInFunctionType)
and arg.output_type() == Lambda
and lamb in arg.args
):
lambda_function_names.add(lamb.value)
return lambda_function_names
def _ensure_lambda_usage_num_input_arguments_valid(
self, prefix: str, name: str, definition: SyntaxTree
):
for function in definition.built_in_functions:
spec = FunctionSpec.from_callable(Functions.get(function.name))
for arg in function.args:
self._ensure_lambda_usage_num_input_arguments_valid(
prefix=prefix, name=name, definition=SyntaxTree([arg])
)
spec = FunctionSpec.from_callable(
name=function.name, callable_ref=Functions.get(function.name)
)
if not (lambda_type := spec.is_lambda_like):
return
lambda_function_names = set(
lamb.value for lamb in SyntaxTree(function.args).lambdas if isinstance(lamb, Lambda)
)
lambda_function_names = self._get_lambda_function_names_to_evaluate(function=function)
# Only case len(lambda_function_names) > 1 is when used in if-statements
for lambda_function_name in lambda_function_names:
if Functions.is_built_in(lambda_function_name):
lambda_spec = FunctionSpec.from_callable(Functions.get(lambda_function_name))
lambda_spec = FunctionSpec.from_callable(
name=lambda_function_name, callable_ref=Functions.get(lambda_function_name)
)
if not lambda_spec.is_num_args_compatible(lambda_type.num_input_args()):
expected_args_str = str(lambda_spec.num_required_args)
if lambda_spec.num_required_args != len(lambda_spec.args):
@ -366,13 +392,15 @@ class Script:
# If the variable's variable dependencies contain an unresolvable variable,
# declare it as unresolvable and continue
elif definition.contains(unresolvable):
elif definition.contains(unresolvable, custom_function_definitions=self._functions):
unresolvable.add(variable)
del unresolved[variable]
# Otherwise, if it has dependencies that are all resolved, then
# resolve the definition
elif not definition.is_subset_of(variables=resolved.keys()):
elif not definition.is_subset_of(
variables=resolved.keys(), custom_function_definitions=self._functions
):
resolved[variable] = unresolved[variable].resolve(
resolved_variables=resolved,
custom_functions=self._functions,

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

@ -117,7 +117,7 @@ class BuiltInFunction(Function, BuiltInFunctionType):
-------
The FunctionSpec of the BuiltInFunction
"""
return FunctionSpec.from_callable(self.callable)
return FunctionSpec.from_callable(name=self.name, callable_ref=self.callable)
@classmethod
def _arg_output_type(cls, arg: Argument) -> Type[Argument]:
@ -229,6 +229,8 @@ class BuiltInFunction(Function, BuiltInFunctionType):
assert isinstance(lambda_array, Array)
if len(lambda_array.value) == 0:
return Array(value=[])
if len(lambda_array.value) == 1:
return lambda_array.value[0]
@ -257,14 +259,27 @@ class BuiltInFunction(Function, BuiltInFunctionType):
resolved_variables: Dict[Variable, Resolvable],
custom_functions: Dict[str, "VariableDependency"],
) -> Resolvable:
# TODO: Make conditionals not execute all branches!!!
conditional_return_args = self.function_spec.conditional_arg_indices(
num_input_args=len(self.args)
)
# Resolve all non-lambda arguments
resolved_arguments: List[Resolvable | Lambda] = [
self._resolve_argument_type(
arg=arg,
resolved_variables=resolved_variables,
custom_functions=custom_functions,
resolved_arguments: List[Resolvable | Lambda | ReturnableArgument] = [
(
self._resolve_argument_type(
arg=arg,
resolved_variables=resolved_variables,
custom_functions=custom_functions,
)
if idx not in conditional_return_args
else ReturnableArgument(
value=functools.partial(
self._resolve_argument_type, arg, resolved_variables, custom_functions
)
)
)
for arg in self.args
for idx, arg in enumerate(self.args)
]
# If a lambda is in a function's arg, resolve it differently

35
src/ytdl_sub/script/types/variable_dependency.py
View file

@ -34,17 +34,21 @@ class VariableDependency(ABC):
Any arguments in the VariableDependency that may or may not need to be resolved.
"""
def _recurse_get(self, ttype: Type[TypeT], subclass: bool = False) -> List[TypeT]:
def _recurse_get(
self, ttype: Type[TypeT], subclass: bool = False, instance: bool = True
) -> List[TypeT]:
output: List[TypeT] = []
for arg in self._iterable_arguments:
if subclass and issubclass(type(arg), ttype):
output.append(arg)
elif isinstance(arg, ttype):
elif instance and isinstance(arg, ttype):
output.append(arg)
elif type(arg) == ttype: # pylint: disable=unidiomatic-typecheck
output.append(arg)
if isinstance(arg, VariableDependency):
# pylint: disable=protected-access
output.extend(arg._recurse_get(ttype))
output.extend(arg._recurse_get(ttype, subclass=subclass, instance=instance))
# pylint: enable=protected-access
return output
@ -57,7 +61,7 @@ class VariableDependency(ABC):
-------
All Variables that this depends on.
"""
return set(self._recurse_get(Variable))
return set(self._recurse_get(Variable, instance=False))
@final
@property
@ -156,19 +160,38 @@ class VariableDependency(ABC):
raise UNREACHABLE
@final
def is_subset_of(self, variables: Iterable[Variable]) -> bool:
def is_subset_of(
self,
variables: Iterable[Variable],
custom_function_definitions: Dict[str, "VariableDependency"],
) -> bool:
"""
Returns
-------
True if it contains all input variables as a dependency. False otherwise.
"""
for custom_function in self.custom_functions:
if custom_function_definitions[custom_function.name].is_subset_of(
variables=variables, custom_function_definitions=custom_function_definitions
):
return True
return not self.variables.issubset(variables)
@final
def contains(self, variables: Iterable[Variable]) -> bool:
def contains(
self,
variables: Iterable[Variable],
custom_function_definitions: Dict[str, "VariableDependency"],
) -> bool:
"""
Returns
-------
True if it contains any of the input variables. False otherwise.
"""
for custom_function in self.custom_functions:
if custom_function_definitions[custom_function.name].contains(
variables=variables, custom_function_definitions=custom_function_definitions
):
return True
return len(self.variables.intersection(variables)) > 0

57
src/ytdl_sub/script/utils/type_checking.py
View file

@ -53,6 +53,25 @@ def get_optional_type(optional_type: Type) -> Type[NamedType]:
return [arg for arg in optional_type.__args__ if arg != type(None)][0]
def _is_union_compatible(
arg_type: Type[NamedType],
expected_union_type: Type[Resolvable | Optional[Resolvable]],
) -> bool:
if issubclass(arg_type, (NamedCustomFunction, Variable)):
return True # custom-function/variable can be anything, so pass for now
# if the input arg is a union, do a direct comparison
if is_union(arg_type):
return arg_type == expected_union_type
# otherwise, iterate the union to see if it's compatible
for union_type in expected_union_type.__args__:
if issubclass(arg_type, union_type):
return True
return False
def _is_type_compatible(
arg_type: Type[NamedType],
expected_arg_type: Type[Resolvable | Optional[Resolvable]],
@ -63,24 +82,11 @@ def _is_type_compatible(
True if arg is compatible with expected_arg_type. False otherwise.
"""
if is_union(expected_arg_type):
# See if the arg is a valid against the union
valid_type = False
return _is_union_compatible(arg_type=arg_type, expected_union_type=expected_arg_type)
# if the input arg is a union, do a direct comparison
if is_union(arg_type):
valid_type = arg_type == expected_arg_type
# otherwise, iterate the union to see if it's compatible
else:
for union_type in expected_arg_type.__args__:
if issubclass(arg_type, union_type):
valid_type = True
break
if not valid_type:
return False
# If the input is a union and the expected type is not, see if
# each possible union input is compatible with the expected type
elif is_union(arg_type):
if is_union(arg_type):
for union_type in arg_type.__args__:
if not _is_type_compatible(union_type, expected_arg_type):
return False
@ -118,6 +124,7 @@ def is_type_compatible(
@dataclass(frozen=True)
class FunctionSpec:
function_name: str
return_type: Type[Resolvable]
arg_names: List[str]
args: Optional[List[Type[Resolvable | Optional[Resolvable]]]] = None
@ -187,6 +194,22 @@ class FunctionSpec:
return sum(1 for arg in self.args if not is_optional(arg))
return 0 # varargs can take any number
def conditional_arg_indices(self, num_input_args: int) -> List[int]:
"""
Returns
-------
If the function is conditional, return the indices of the arguments that
return for different branches.
"""
if self.function_name == "if":
return [1, 2] # true, false
if self.function_name == "elif":
# if, retA, elif, retB, retElse
return list(range(1, num_input_args, 2)) + [num_input_args - 1]
if self.function_name == "if_passthrough":
return [0, 1] # true-passthrough, false-passthrough
return []
@property
def is_lambda_reduce_function(self) -> Optional[Type[LambdaReduce]]:
"""
@ -261,7 +284,7 @@ class FunctionSpec:
return self._to_human_readable_name(self.return_type)
@classmethod
def from_callable(cls, callable_ref: Callable[..., Resolvable]) -> "FunctionSpec":
def from_callable(cls, name: str, callable_ref: Callable[..., Resolvable]) -> "FunctionSpec":
"""
Returns
-------
@ -270,12 +293,14 @@ class FunctionSpec:
arg_spec: FullArgSpec = inspect.getfullargspec(callable_ref)
if arg_spec.varargs:
return FunctionSpec(
function_name=name,
return_type=arg_spec.annotations["return"],
arg_names=[arg_spec.varargs],
varargs=arg_spec.annotations[arg_spec.varargs],
)
return FunctionSpec(
function_name=name,
return_type=arg_spec.annotations["return"],
arg_names=arg_spec.args,
args=[arg_spec.annotations[arg_name] for arg_name in arg_spec.args],

3
src/ytdl_sub/utils/script.py
View file

@ -13,6 +13,7 @@ from ytdl_sub.script.types.resolvable import Argument
from ytdl_sub.script.types.resolvable import Boolean
from ytdl_sub.script.types.resolvable import Float
from ytdl_sub.script.types.resolvable import Integer
from ytdl_sub.script.types.resolvable import Lambda
from ytdl_sub.script.types.resolvable import String
from ytdl_sub.script.types.variable import Variable
from ytdl_sub.script.utils.exceptions import UNREACHABLE
@ -115,6 +116,8 @@ class ScriptUtils:
out = arg.name
elif isinstance(arg, Function):
out = f"%{arg.name}( {', '.join(cls._to_script_code(val) for val in arg.args)} )"
elif isinstance(arg, Lambda):
out = f"%{arg.value}"
else:
raise UNREACHABLE
return f"{{ {out} }}" if top_level else out

34
tests/unit/script/functions/test_conditional_functions.py
View file

@ -101,3 +101,37 @@ class TestConditionalFunction:
)
}"""
)
@pytest.mark.parametrize(
"function_str, expected_output",
[
("{%if(True, True, %assert(False, 'should not reach'))}", True),
("{%if(False, %assert(False, 'should not reach'), False)}", False),
],
)
def test_if_function_only_evaluates_branch(self, function_str: str, expected_output: bool):
output = single_variable_output(function_str)
assert output == expected_output
@pytest.mark.parametrize(
"function_str, expected_output",
[
("{%elif(True, True, %assert(False, 'should not reach'))}", True),
("{%elif(False, %assert(False, 'should not reach'), False)}", False),
],
)
def test_elif_function_only_evaluates_branch(self, function_str: str, expected_output: bool):
output = single_variable_output(function_str)
assert output == expected_output
@pytest.mark.parametrize(
"function_str, expected_output",
[
("{%if_passthrough(True, %assert(False, 'should not reach'))}", True),
],
)
def test_if_passthrough_function_only_evaluates_branch(
self, function_str: str, expected_output: bool
):
output = single_variable_output(function_str)
assert output == expected_output

22
tests/unit/script/types/test_lambda_function.py
View file

@ -15,6 +15,11 @@ class TestLambdaFunction:
{"%times_two": "{%mul($0, 2)}", "wip": "{%array_apply([1, 2, 3], %times_two)}"}
).resolve() == ScriptOutput({"wip": Array([Integer(2), Integer(4), Integer(6)])})
def test_lambda_with_custom_function_empty_input(self):
assert Script(
{"%times_two": "{%mul($0, 2)}", "wip": "{%array_apply([], %times_two)}"}
).resolve() == ScriptOutput({"wip": Array([])})
def test_conditional_lambda_with_custom_functions(self):
assert Script(
{
@ -54,6 +59,23 @@ class TestLambdaFunction:
}
).resolve() == ScriptOutput({"output": Integer(4)})
def test_multiple_lambdas_single_definition(self):
url_map_def = """{
%array_reduce(
%array_apply( array_def, %array_map_format),
%map_extend
)
}"""
script = Script(
{
"%array_map_format": "{ {$0: $0 } }",
"array_def": "{ [1, 2, 3] }",
"category_url_map": url_map_def,
}
)
assert script.resolve().get("category_url_map").native == {1: 1, 2: 2, 3: 3}
class TestLambdaFunctionIncompatibleNumArguments:
@pytest.mark.parametrize(

2
tools/docgen/scripting_functions.py
View file

@ -30,7 +30,7 @@ def function_class_to_name(obj: Type[Any]) -> str:
def function_type_hinting(display_function_name: str, function: Any) -> str:
spec = FunctionSpec.from_callable(function)
spec = FunctionSpec.from_callable(name=display_function_name, callable_ref=function)
out = ":spec: ``"
out += display_function_name
out += spec.human_readable_input_args()