Add Support for Function Builders

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"]

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]]
 
 

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"""