From 99d38b55fb57364ad9520efaf270043f75c9ec22 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Thu, 9 Nov 2023 17:02:03 -0800 Subject: [PATCH] custom functions --- src/ytdl_sub/script/parser.py | 30 ++++- src/ytdl_sub/script/script.py | 36 ++++++ src/ytdl_sub/script/syntax_tree.py | 42 +++++- src/ytdl_sub/script/types/array.py | 33 ++++- src/ytdl_sub/script/types/function.py | 120 ++++++++++++++---- src/ytdl_sub/script/types/map.py | 33 ++++- src/ytdl_sub/script/types/variable.py | 5 + .../script/types/variable_dependency.py | 21 ++- tests/unit/script/test_parser.py | 4 + tests/unit/script/test_script.py | 37 ++++++ 10 files changed, 317 insertions(+), 44 deletions(-) create mode 100644 src/ytdl_sub/script/script.py create mode 100644 tests/unit/script/test_script.py diff --git a/src/ytdl_sub/script/parser.py b/src/ytdl_sub/script/parser.py index 89081b63..566004a4 100644 --- a/src/ytdl_sub/script/parser.py +++ b/src/ytdl_sub/script/parser.py @@ -12,6 +12,7 @@ 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 String +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.utils.exceptions import StringFormattingException from ytdl_sub.validators.string_formatter_validators import is_valid_source_variable_name @@ -71,6 +72,30 @@ class _Parser: assert is_valid_source_variable_name(var_name, raise_exception=False) return Variable(var_name) + def _parse_function_argument(self) -> FunctionArgument: + """ + Begin parsing function args after the first ``$``, i.e. ``$1`` + """ + var_name = "" + while ch := self._read(increment_pos=False): + if ch.isspace() and not var_name: + self._pos += 1 + continue + if ch in ["}", ",", ")", "]"] or ch.isspace(): + break + + is_numeric = ch.isnumeric() + if not is_numeric: + raise StringFormattingException("invalid function var name") + + var_name += ch + self._pos += 1 + + if not var_name: + raise StringFormattingException("invalid var name") + + return FunctionArgument(name=f"${var_name}") + def _parse_numeric(self) -> Integer | Float: numeric_string = "" while ch := self._read(increment_pos=False): @@ -125,6 +150,9 @@ class _Parser: if self._read(increment_pos=False) == "{": self._pos += 1 return self._parse_map() + if self._read(increment_pos=False) == "$": + self._pos += 1 + return self._parse_function_argument() if self._read(increment_pos=False).isascii() and self._read(increment_pos=False).islower(): return self._parse_variable() raise StringFormattingException( @@ -167,7 +195,7 @@ class _Parser: while ch := self._read(): if ch == ")": - return Function(name=function_name, args=function_args) + return Function.from_name_and_args(name=function_name, args=function_args) if ch != "(": function_name += ch diff --git a/src/ytdl_sub/script/script.py b/src/ytdl_sub/script/script.py new file mode 100644 index 00000000..d7e3bd37 --- /dev/null +++ b/src/ytdl_sub/script/script.py @@ -0,0 +1,36 @@ +from typing import Dict, Optional + +from ytdl_sub.script.parser import parse +from ytdl_sub.script.syntax_tree import SyntaxTree +from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.variable import Variable + + +class Script: + @classmethod + def _is_function(cls, override_name: str): + return override_name.startswith("%") + + @classmethod + def _function_name(self, function_key: str) -> str: + return function_key[1:] + + def __init__(self, overrides: Dict[str, str]): + self._functions: Dict[str, SyntaxTree] = { + self._function_name(function_key): parse(function_value) + for function_key, function_value in overrides.items() + if self._is_function(function_key) + } + + self._variables: Dict[str, SyntaxTree] = { + override_name: parse(override_value) + for override_name, override_value in overrides.items() + if not self._is_function(override_name) + } + + def resolve(self, pre_resolved_variables: Optional[Dict[Variable, Resolvable]] = None) -> Dict[str, Resolvable]: + return SyntaxTree.resolve_overrides( + parsed_overrides=self._variables, + custom_functions=self._functions, + pre_resolved_variables=pre_resolved_variables, + ) \ No newline at end of file diff --git a/src/ytdl_sub/script/syntax_tree.py b/src/ytdl_sub/script/syntax_tree.py index 317d701a..666594c4 100644 --- a/src/ytdl_sub/script/syntax_tree.py +++ b/src/ytdl_sub/script/syntax_tree.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Dict +from typing import Dict, Optional from typing import List from typing import Set @@ -7,6 +7,7 @@ from ytdl_sub.script.types.function import Function from ytdl_sub.script.types.resolvable import ArgumentType from ytdl_sub.script.types.resolvable import Resolvable from ytdl_sub.script.types.resolvable import String +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.script.types.variable_dependency import VariableDependency from ytdl_sub.utils.exceptions import StringFormattingException @@ -27,16 +28,40 @@ class SyntaxTree(VariableDependency): for token in self.ast: if isinstance(token, Variable): variables.add(token) - elif isinstance(token, Function): + elif isinstance(token, VariableDependency): variables.update(token.variables) return variables - def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: + @property + def function_arguments(self) -> Set[FunctionArgument]: + """ + Returns + ------- + All function arguments used within the SyntaxTree + """ + function_arguments: Set[FunctionArgument] = set() + for token in self.ast: + if isinstance(token, FunctionArgument): + function_arguments.add(token) + elif isinstance(token, VariableDependency): + function_arguments.update(token.function_arguments) + + return function_arguments + + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], + ) -> Resolvable: resolved: List[Resolvable] = [] for token in self.ast: resolved.append( - self._resolve_argument_type(resolved_variables=resolved_variables, arg=token) + self._resolve_argument_type( + arg=token, + resolved_variables=resolved_variables, + custom_functions=custom_functions, + ) ) # If only one resolvable resides in the AST, return as that @@ -47,13 +72,15 @@ class SyntaxTree(VariableDependency): return String("".join([str(res) for res in resolved])) @classmethod - def resolve_overrides(cls, parsed_overrides: Dict[str, "SyntaxTree"]) -> Dict[str, Resolvable]: + def resolve_overrides( + cls, parsed_overrides: Dict[str, "SyntaxTree"], custom_functions: Dict[str, "SyntaxTree"], pre_resolved_variables: Optional[Dict[Variable, Resolvable]] + ) -> Dict[str, Resolvable]: overrides: Dict[Variable, "SyntaxTree"] = { Variable(name): ast for name, ast in parsed_overrides.items() } unresolved_variables: List[Variable] = list(overrides.keys()) - resolved_variables: Dict[Variable, Resolvable] = {} + resolved_variables: Dict[Variable, Resolvable] = pre_resolved_variables if pre_resolved_variables else {} while unresolved_variables: unresolved_count: int = len(unresolved_variables) @@ -63,7 +90,8 @@ class SyntaxTree(VariableDependency): resolved_variables=resolved_variables ): resolved_variables[variable] = overrides[variable].resolve( - resolved_variables=resolved_variables + resolved_variables=resolved_variables, + custom_functions=custom_functions, ) unresolved_variables.remove(variable) diff --git a/src/ytdl_sub/script/types/array.py b/src/ytdl_sub/script/types/array.py index a2aaa77c..b11e6672 100644 --- a/src/ytdl_sub/script/types/array.py +++ b/src/ytdl_sub/script/types/array.py @@ -5,6 +5,7 @@ from typing import Set from ytdl_sub.script.types.resolvable import ArgumentType from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.script.types.variable_dependency import VariableDependency @@ -20,12 +21,38 @@ class UnresolvedArray(Array, VariableDependency, ArgumentType): @property def variables(self) -> Set[Variable]: - return {value for value in self.value if isinstance(value, Variable)} + variables: Set[Variable] = set() + for arg in self.value: + if isinstance(arg, Variable): + variables.add(arg) + elif isinstance(arg, VariableDependency): + variables.update(arg.variables) - def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: + return variables + + @property + def function_arguments(self) -> Set[FunctionArgument]: + variables: Set[FunctionArgument] = set() + for arg in self.value: + if isinstance(arg, FunctionArgument): + variables.add(arg) + elif isinstance(arg, VariableDependency): + variables.update(arg.function_arguments) + + return variables + + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], + ) -> Resolvable: return ResolvedArray( [ - self._resolve_argument_type(resolved_variables=resolved_variables, arg=arg) + self._resolve_argument_type( + arg=arg, + resolved_variables=resolved_variables, + custom_functions=custom_functions, + ) for arg in self.value ] ) diff --git a/src/ytdl_sub/script/types/function.py b/src/ytdl_sub/script/types/function.py index c524e01c..e8f8c706 100644 --- a/src/ytdl_sub/script/types/function.py +++ b/src/ytdl_sub/script/types/function.py @@ -1,5 +1,7 @@ +import copy import functools import inspect +from abc import ABC from dataclasses import dataclass from inspect import FullArgSpec from typing import Callable @@ -17,6 +19,7 @@ from ytdl_sub.script.types.resolvable import Resolvable from ytdl_sub.script.types.resolvable import Resolvable_0 from ytdl_sub.script.types.resolvable import Resolvable_1 from ytdl_sub.script.types.resolvable import Resolvable_2 +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.script.types.variable_dependency import VariableDependency from ytdl_sub.utils.exceptions import StringFormattingException @@ -122,21 +125,87 @@ class FunctionInputSpec: @dataclass(frozen=True) -class Function(VariableDependency, ArgumentType): +class Function(VariableDependency, ArgumentType, ABC): name: str args: List[ArgumentType] - def __post_init__(self): - if not self.input_spec.is_compatible(input_args=self.args): - raise StringFormattingException( - f"Invalid arguments passed to function {self.name}.\n" - f"{self._expected_received_error_msg()}" - ) + @property + def variables(self) -> Set[Variable]: + """ + Returns + ------- + All variables used within the function + """ + variables: Set[Variable] = set() + for arg in self.args: + if isinstance(arg, Variable): + variables.add(arg) + elif isinstance(arg, VariableDependency): + variables.update(arg.variables) + return variables + + @property + def function_arguments(self) -> Set[FunctionArgument]: + """ + Returns + ------- + All function arguments used within the function + """ + function_arguments: Set[FunctionArgument] = set() + for arg in self.args: + if isinstance(arg, FunctionArgument): + function_arguments.add(arg) + elif isinstance(arg, VariableDependency): + function_arguments.update(arg.function_arguments) + + return function_arguments + + @classmethod + def from_name_and_args(cls, name: str, args: List[ArgumentType]) -> "Function": + if hasattr(Functions, name) or hasattr(Functions, name + "_"): + return BuiltInFunction(name=name, args=args) + + return CustomFunction(name=name, args=args) + + +class CustomFunction(Function): + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], + ) -> Resolvable: + resolved_args: List[Resolvable] = [ + self._resolve_argument_type( + arg=arg, resolved_variables=resolved_variables, custom_functions=custom_functions + ) + for arg in self.args + ] + + if self.name in custom_functions: + if len(self.args) != len(custom_functions[self.name].function_arguments): + raise StringFormattingException("Custom function arg length does not equal") + + resolved_variables_with_args = copy.deepcopy(resolved_variables) + for i, arg in enumerate(resolved_args): + function_arg = FunctionArgument(name=f"${i+1}") # Function args are 1-based + if function_arg in resolved_variables_with_args: + raise StringFormattingException("nested custom functions???") + resolved_variables_with_args[function_arg] = arg + + return custom_functions[self.name].resolve( + resolved_variables=resolved_variables_with_args, + custom_functions=custom_functions, + ) + else: + raise StringFormattingException(f"Custom function {self.name} does not exist") + + +class BuiltInFunction(Function): def _expected_received_error_msg(self) -> str: received_type_names: List[str] = [] for arg in self.args: - if isinstance(arg, Function): + if isinstance(arg, BuiltInFunction): received_type_names.append(f"%{arg.name}(...)->{arg.output_type.__name__}") else: received_type_names.append(arg.__class__.__name__) @@ -145,6 +214,13 @@ class Function(VariableDependency, ArgumentType): return f"Expected {self.input_spec.expected_args_str()}.\nReceived {received_args_str}" + def __post_init__(self): + if not self.input_spec.is_compatible(input_args=self.args): + raise StringFormattingException( + f"Invalid arguments passed to function {self.name}.\n" + f"{self._expected_received_error_msg()}" + ) + @property def callable(self) -> Callable[..., Resolvable]: if hasattr(Functions, self.name): @@ -181,25 +257,15 @@ class Function(VariableDependency, ArgumentType): return output_type - @property - def variables(self) -> Set[Variable]: - """ - Returns - ------- - All variables used within the function - """ - variables: Set[Variable] = set() - for arg in self.args: - if isinstance(arg, Variable): - variables.add(arg) - elif isinstance(arg, Function): - variables.update(arg.variables) - - return variables - - def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: - resolved_args = [ - self._resolve_argument_type(resolved_variables=resolved_variables, arg=arg) + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], + ) -> Resolvable: + resolved_args: List[Resolvable] = [ + self._resolve_argument_type( + arg=arg, resolved_variables=resolved_variables, custom_functions=custom_functions + ) for arg in self.args ] diff --git a/src/ytdl_sub/script/types/map.py b/src/ytdl_sub/script/types/map.py index 21a2016d..a76976da 100644 --- a/src/ytdl_sub/script/types/map.py +++ b/src/ytdl_sub/script/types/map.py @@ -6,6 +6,7 @@ from typing import Set from ytdl_sub.script.types.resolvable import ArgumentType from ytdl_sub.script.types.resolvable import Hashable from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.script.types.variable_dependency import VariableDependency from ytdl_sub.utils.exceptions import StringFormattingException @@ -26,21 +27,47 @@ class UnresolvedMap(Map, VariableDependency, ArgumentType): for key, value in self.value.items(): if isinstance(key, Variable): output.add(key) + elif isinstance(key, VariableDependency): + output.update(key.variables) + if isinstance(value, Variable): output.add(key) + elif isinstance(value, VariableDependency): + output.update(value.variables) + return output - def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: + @property + def function_arguments(self) -> Set[FunctionArgument]: + output: Set[FunctionArgument] = set() + for key, value in self.value.items(): + if isinstance(key, FunctionArgument): + output.add(key) + elif isinstance(key, VariableDependency): + output.update(key.function_arguments) + + if isinstance(value, FunctionArgument): + output.add(key) + elif isinstance(value, VariableDependency): + output.update(value.function_arguments) + + return output + + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, VariableDependency], + ) -> Resolvable: output: Dict[Hashable, Resolvable] = {} for key, value in self.value.items(): resolved_key = self._resolve_argument_type( - resolved_variables=resolved_variables, arg=key + arg=key, resolved_variables=resolved_variables, custom_functions=custom_functions ) if not isinstance(resolved_key, Hashable): raise StringFormattingException("key is not hashable") output[resolved_key] = self._resolve_argument_type( - resolved_variables=resolved_variables, arg=value + arg=value, resolved_variables=resolved_variables, custom_functions=custom_functions ) return ResolvedMap(output) diff --git a/src/ytdl_sub/script/types/variable.py b/src/ytdl_sub/script/types/variable.py index af723724..7af9d36f 100644 --- a/src/ytdl_sub/script/types/variable.py +++ b/src/ytdl_sub/script/types/variable.py @@ -6,3 +6,8 @@ from ytdl_sub.script.types.resolvable import ArgumentType @dataclass(frozen=True) class Variable(ArgumentType): name: str + + +@dataclass(frozen=True) +class FunctionArgument(Variable): + pass diff --git a/src/ytdl_sub/script/types/variable_dependency.py b/src/ytdl_sub/script/types/variable_dependency.py index b08f821b..66e9f337 100644 --- a/src/ytdl_sub/script/types/variable_dependency.py +++ b/src/ytdl_sub/script/types/variable_dependency.py @@ -7,6 +7,7 @@ from typing import final from ytdl_sub.script.types.resolvable import ArgumentType from ytdl_sub.script.types.resolvable import Resolvable +from ytdl_sub.script.types.variable import FunctionArgument from ytdl_sub.script.types.variable import Variable from ytdl_sub.utils.exceptions import StringFormattingException @@ -18,12 +19,24 @@ class VariableDependency(ABC): def variables(self) -> Set[Variable]: raise NotImplemented() + @property @abstractmethod - def resolve(self, resolved_variables: Dict[Variable, Resolvable]) -> Resolvable: + def function_arguments(self) -> Set[FunctionArgument]: + raise NotImplemented() + + @abstractmethod + def resolve( + self, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], + ) -> Resolvable: raise NotImplemented() def _resolve_argument_type( - self, resolved_variables: Dict[Variable, Resolvable], arg: ArgumentType + self, + arg: ArgumentType, + resolved_variables: Dict[Variable, Resolvable], + custom_functions: Dict[str, "VariableDependency"], ) -> Resolvable: if isinstance(arg, Resolvable): return arg @@ -32,7 +45,9 @@ class VariableDependency(ABC): raise StringFormattingException("should never reach@") return resolved_variables[arg] if isinstance(arg, VariableDependency): - return arg.resolve(resolved_variables) + return arg.resolve( + resolved_variables=resolved_variables, custom_functions=custom_functions + ) assert False, "never reach here" diff --git a/tests/unit/script/test_parser.py b/tests/unit/script/test_parser.py index 896b3979..aacab1a6 100644 --- a/tests/unit/script/test_parser.py +++ b/tests/unit/script/test_parser.py @@ -47,6 +47,10 @@ class TestParser: parsed_extend.resolve({}) assert False + def test_function_argument(self): + parsed = parse("hello {%map([$1, $2])}") + assert False + def test_conditional(self): parsed = parse("hello {%if(True, 'hi', 3.4)}") assert parsed == SyntaxTree( diff --git a/tests/unit/script/test_script.py b/tests/unit/script/test_script.py new file mode 100644 index 00000000..378c18d3 --- /dev/null +++ b/tests/unit/script/test_script.py @@ -0,0 +1,37 @@ +from typing import Dict + +import pytest + +from ytdl_sub.script.script import Script +from ytdl_sub.script.syntax_tree import SyntaxTree +from ytdl_sub.script.types.function import Function +from ytdl_sub.script.types.resolvable import String +from ytdl_sub.script.types.variable import Variable +from ytdl_sub.utils.exceptions import StringFormattingException + + +class TestSyntaxTree: + def test_simple(self): + script = Script( + { + "a": "a", + "b": "{b_}", + "b_": "b", + } + ) + + def test_custom_function(self): + script = Script( + { + "%custom_func": "return {[$1, $2]}", + "aa": "a", + "bb": "b", + "cc": "{%custom_func(aa, bb)}", + } + ) + + out = script.resolve() + assert False + + +