266 lines
9.5 KiB
Python
266 lines
9.5 KiB
Python
import json
|
|
import re
|
|
from typing import Any, Dict, Optional
|
|
|
|
from ytdl_sub.entries.script.custom_functions import CustomFunctions
|
|
from ytdl_sub.entries.script.variable_definitions import VARIABLES
|
|
from ytdl_sub.entries.script.variable_types import BooleanVariable, IntegerVariable
|
|
from ytdl_sub.script.parser import parse
|
|
from ytdl_sub.script.types.array import Array, UnresolvedArray
|
|
from ytdl_sub.script.types.function import BuiltInFunction, Function
|
|
from ytdl_sub.script.types.map import Map, UnresolvedMap
|
|
from ytdl_sub.script.types.resolvable import Argument, Boolean, Float, Integer, Lambda, String
|
|
from ytdl_sub.script.types.syntax_tree import SyntaxTree
|
|
from ytdl_sub.script.types.variable import Variable
|
|
from ytdl_sub.script.utils.exceptions import UNREACHABLE
|
|
from ytdl_sub.script.utils.name_validation import is_function
|
|
|
|
# pylint: disable=too-many-return-statements
|
|
# pylint: disable=too-many-branches
|
|
|
|
|
|
class ScriptUtils:
|
|
@classmethod
|
|
def add_sanitized_variables(cls, variables: Dict[str, str]) -> Dict[str, str]:
|
|
"""
|
|
Helper to add sanitized variables to a Script
|
|
"""
|
|
sanitized_variables = {
|
|
f"{name}_sanitized": f"{{%sanitize({name})}}"
|
|
for name in variables.keys()
|
|
if not is_function(name)
|
|
}
|
|
return dict(variables, **sanitized_variables)
|
|
|
|
@classmethod
|
|
def add_sanitized_parsed_variables(
|
|
cls, variables: Dict[str, SyntaxTree]
|
|
) -> Dict[str, SyntaxTree]:
|
|
"""
|
|
Helper to add sanitized variables to a Script
|
|
"""
|
|
sanitized_variables = {
|
|
f"{name}_sanitized": SyntaxTree(
|
|
ast=[BuiltInFunction(name="sanitize", args=[Variable(name)])]
|
|
)
|
|
for name in variables.keys()
|
|
if not is_function(name)
|
|
}
|
|
return variables | sanitized_variables
|
|
|
|
@classmethod
|
|
def to_script(cls, value: Any, sort_keys: bool = True) -> str:
|
|
"""
|
|
Converts a python value to a script value
|
|
"""
|
|
if value is None:
|
|
out = ""
|
|
elif isinstance(value, str):
|
|
out = value
|
|
elif isinstance(value, bool):
|
|
out = f"{{%bool({value})}}"
|
|
elif isinstance(value, int):
|
|
out = f"{{%int({value})}}"
|
|
elif isinstance(value, float):
|
|
out = f"{{%float({value})}}"
|
|
else:
|
|
dumped_json = json.dumps(value, ensure_ascii=False, sort_keys=sort_keys)
|
|
# Remove triple-single-quotes from JSON to avoid parsing issues
|
|
dumped_json = re.sub("'{3,}", "'", dumped_json)
|
|
|
|
out = f"{{%from_json('''{dumped_json}''')}}"
|
|
|
|
return out
|
|
|
|
@classmethod
|
|
def _to_script_argument(cls, value: Any) -> Argument:
|
|
# Handle simple types as above
|
|
if value is None or (isinstance(value, str) and value == ""):
|
|
return String("")
|
|
if isinstance(value, str):
|
|
ast = parse(text=value).ast
|
|
if len(ast) == 1:
|
|
return ast[0]
|
|
|
|
return BuiltInFunction(name="concat", args=ast)
|
|
if isinstance(value, bool):
|
|
return Boolean(value)
|
|
if isinstance(value, int):
|
|
return Integer(value)
|
|
if isinstance(value, float):
|
|
return Float(value)
|
|
if isinstance(value, list):
|
|
return UnresolvedArray([cls._to_script_argument(val) for val in value])
|
|
if isinstance(value, dict):
|
|
return UnresolvedMap(
|
|
{
|
|
cls._to_script_argument(key): cls._to_script_argument(val)
|
|
for key, val in value.items()
|
|
}
|
|
)
|
|
|
|
raise UNREACHABLE
|
|
|
|
@classmethod
|
|
def _get_quote_char(cls, arg: str) -> str:
|
|
contains_single_quote = "'" in arg
|
|
contains_double_quote = '"' in arg
|
|
|
|
if not contains_single_quote and not contains_double_quote:
|
|
return '"'
|
|
if not contains_single_quote and contains_double_quote:
|
|
return "'"
|
|
if contains_single_quote and not contains_double_quote:
|
|
return '"'
|
|
return "'''"
|
|
|
|
@classmethod
|
|
def _maybe_to_optimized_sanitize(cls, arg: Argument) -> Argument:
|
|
# If it is %sanitize(%concat(...)), return %sanitize(...)
|
|
if (
|
|
isinstance(arg, Function)
|
|
and arg.name == "sanitize"
|
|
and len(arg.args) == 1
|
|
and isinstance(arg.args[0], Function)
|
|
and arg.args[0].name == "concat"
|
|
):
|
|
return BuiltInFunction(name="sanitize", args=arg.args[0].args)
|
|
|
|
return arg
|
|
|
|
@classmethod
|
|
def _maybe_sanitized_script_code(cls, arg: Argument) -> Optional[str]:
|
|
if not (isinstance(arg, Function) and arg.name == "sanitize"):
|
|
return None
|
|
|
|
output = ""
|
|
for sub_arg in arg.args:
|
|
if isinstance(sub_arg, Variable):
|
|
# No need to sanitize built-in integer variables
|
|
if isinstance(VARIABLES.get(sub_arg.name), (IntegerVariable, BooleanVariable)):
|
|
output += f"{{ {sub_arg.name} }}"
|
|
else:
|
|
output += f"{{ {sub_arg.name}_sanitized }}"
|
|
elif isinstance(sub_arg, (Integer, Float, Boolean)):
|
|
output += str(sub_arg.native)
|
|
elif isinstance(sub_arg, String):
|
|
output += CustomFunctions.sanitize(sub_arg).native
|
|
elif isinstance(sub_arg, BuiltInFunction) and (
|
|
issubclass(sub_arg.function_spec.return_type, (Integer, Float, Boolean))
|
|
or sub_arg.name == "pad_zero"
|
|
):
|
|
# If we know the function's output is sanitized, let's not wrap it
|
|
output += cls._to_script_code(sub_arg, top_level=True)
|
|
else:
|
|
# Purposefully do not set top_level to True so we do not recurse
|
|
output += (
|
|
f"{{ {cls._to_script_code(BuiltInFunction(name='sanitize', args=[sub_arg]))} }}"
|
|
)
|
|
|
|
return output
|
|
|
|
@classmethod
|
|
def _maybe_concat_script_code(cls, arg: Argument) -> Optional[str]:
|
|
if not (isinstance(arg, Function) and arg.name == "concat"):
|
|
return None
|
|
|
|
out = ""
|
|
for sub_arg in arg.args:
|
|
out += cls._to_script_code(sub_arg, top_level=True)
|
|
|
|
return out
|
|
|
|
@classmethod
|
|
def _to_script_code(cls, arg: Argument, top_level: bool = False) -> str:
|
|
if not top_level and isinstance(arg, (Integer, Boolean, Float)):
|
|
return str(arg.native)
|
|
|
|
if isinstance(arg, String):
|
|
if arg.native == "":
|
|
return "" if top_level else "''"
|
|
|
|
quote = cls._get_quote_char(arg.native)
|
|
|
|
return arg.native if top_level else f"{quote}{arg.native}{quote}"
|
|
|
|
arg = cls._maybe_to_optimized_sanitize(arg)
|
|
|
|
if top_level:
|
|
if (out := cls._maybe_sanitized_script_code(arg)) is not None:
|
|
return out
|
|
if (out := cls._maybe_concat_script_code(arg)) is not None:
|
|
return out
|
|
|
|
if isinstance(arg, Integer):
|
|
out = f"%int({arg.native})"
|
|
elif isinstance(arg, Boolean):
|
|
out = f"%bool({arg.native})"
|
|
elif isinstance(arg, Float):
|
|
out = f"%float({arg.native})"
|
|
elif isinstance(arg, (Array, UnresolvedArray)):
|
|
out = f"[ {', '.join(cls._to_script_code(val) for val in arg.value)} ]"
|
|
elif isinstance(arg, (Map, UnresolvedMap)):
|
|
kv_list = (
|
|
f"{cls._to_script_code(key)}: {cls._to_script_code(val)}"
|
|
for key, val in arg.value.items()
|
|
)
|
|
out = f"{{ {', '.join(kv_list)} }}"
|
|
elif isinstance(arg, Variable):
|
|
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
|
|
|
|
@classmethod
|
|
def _is_top_level_string(cls, tree: SyntaxTree) -> Optional[str]:
|
|
if not (
|
|
len(tree.ast) == 1
|
|
and isinstance(tree.ast[0], BuiltInFunction)
|
|
and tree.ast[0].name == "concat"
|
|
):
|
|
return None
|
|
|
|
output = ""
|
|
for arg in tree.ast[0].args:
|
|
if isinstance(arg, BuiltInFunction) and arg.name == "string" and len(arg.args) == 1:
|
|
output += cls._to_script_code(arg.args[0], top_level=True)
|
|
else:
|
|
output += cls._to_script_code(arg, top_level=True)
|
|
|
|
return output
|
|
|
|
@classmethod
|
|
def _syntax_tree_to_native_script(cls, tree: SyntaxTree) -> str:
|
|
if (output := cls._is_top_level_string(tree)) is not None:
|
|
return output
|
|
|
|
output = ""
|
|
for arg in tree.ast:
|
|
output += cls._to_script_code(arg, top_level=True)
|
|
return output
|
|
|
|
@classmethod
|
|
def to_native_script(cls, value: Any) -> str:
|
|
"""
|
|
Converts any JSON-compatible value into equivalent script syntax
|
|
"""
|
|
if isinstance(value, SyntaxTree):
|
|
return cls._syntax_tree_to_native_script(value)
|
|
|
|
return cls._to_script_code(cls._to_script_argument(value), top_level=True)
|
|
|
|
@classmethod
|
|
def bool_formatter_output(cls, output: str) -> bool:
|
|
"""
|
|
Translate formatter output to a boolean
|
|
"""
|
|
if not output or output.lower() == "false":
|
|
return False
|
|
try:
|
|
return bool(json.loads(output))
|
|
except Exception: # pylint: disable=broad-except
|
|
return True
|