diff --git a/jsonquerylang/compile.py b/jsonquerylang/compile.py index 3e2d796..de759a4 100644 --- a/jsonquerylang/compile.py +++ b/jsonquerylang/compile.py @@ -1,9 +1,13 @@ import json -from typing import Callable, Optional, Final +from typing import Any, Callable, Mapping, Optional, Final, cast from jsonquerylang.functions import get_functions from jsonquerylang.types import JsonQueryType, JsonType, JsonQueryOptions +CompileFunc = Callable[..., Callable[[Any], Any]] +FunctionDict = Mapping[str, Callable[..., Any]] + + def compile( query: JsonQueryType, options: Optional[JsonQueryOptions] = None ) -> Callable[[JsonType], JsonType]: @@ -33,10 +37,38 @@ def compile( :return: Returns a function which can execute the query """ - functions = get_functions(lambda q: compile(q, options), build_function) + if options is None: + options = {} + + compile_inner: CompileFunc = lambda query: compile(query, options) + + builtin_functions: Final[FunctionDict] = get_functions( + compile_inner, build_function + ) + + custom_functions: Final[FunctionDict] = cast( + FunctionDict, options.get("functions", {}) + ) + + custom_function_builders: Final[list[Callable[[CompileFunc], FunctionDict]]] = ( + options.get("function_builders", []) + ) + + built_custom_function_dicts = [ + function_builder(compile_inner) for function_builder in custom_function_builders + ] + + function_dicts: list[FunctionDict] = [ + builtin_functions, + custom_functions, + *built_custom_function_dicts, + ] - custom_functions: Final = (options.get("functions") if options else None) or {} - all_functions: Final = {**functions, **custom_functions} + all_functions = { + function_name: function_instance + for function_dict in function_dicts + for function_name, function_instance in function_dict.items() + } if type(query) is list: # a function like ["sort", ["get", "name"], "desc"] diff --git a/jsonquerylang/types.py b/jsonquerylang/types.py index dd47ed5..2a31b3e 100644 --- a/jsonquerylang/types.py +++ b/jsonquerylang/types.py @@ -1,4 +1,4 @@ -from typing import TypeAlias, List, Mapping, TypedDict, Callable, NotRequired +from typing import Any, TypeAlias, List, Mapping, TypedDict, Callable, NotRequired JsonType: TypeAlias = ( List["JsonValueType"] | Mapping[str, "JsonValueType"] | "JsonValueType" @@ -46,7 +46,7 @@ class CustomOperatorAfter(TypedDict): class JsonQueryOptions(TypedDict): - functions: NotRequired[Mapping[str, Callable]] + functions: NotRequired[Mapping[str, Callable[..., Any]]] operators: NotRequired[list[CustomOperator]] diff --git a/tests/test_compile.py b/tests/test_compile.py index 9db0d91..7b87790 100644 --- a/tests/test_compile.py +++ b/tests/test_compile.py @@ -101,6 +101,27 @@ def about_eq(a, b): self.assertTrue(go({"a": 2}, ["aboutEq", ["get", "a"], 2], options)) self.assertTrue(go({"a": 1.999}, ["aboutEq", ["get", "a"], 2], options)) + def test_options_function_builders(self): + """should get custom functions from function builders""" + + def sample_function_builder(compile): + def fn_first(values): + compiled_value = compile(values) + + return lambda data: compiled_value(data)[0] + + return { + "first": fn_first, + } + + options = {"function_builders": [sample_function_builder]} + + result = go( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], ["first", ["first", ["get"]]], options + ) + + self.assertEqual(result, 1) + def test_error_handling1(self): """should throw a helpful error when a pipe contains a compile time error"""