Repository: jmespath/jmespath.py

Issue:
Custom functions
What are your thoughts about adding a method to register custom functions directly into `RuntimeFunctions` class in `functions.py`?

JMESPath is almost good enough to use as a domain specific language for general language transforming objects. You can sneak in literals in the multi-select hash. You can filter for values, and transform them to booleans using `<=`, `==` operators. There's some support for making sure values are numbers.

However I don't see anyway to do something like "if value is x, return `y`" where you show a literal if a condition matches. There's no way to convert conditions to an arbitrary literal - if a value in a multi-select hash is going to be a literal, it has to be of the same value no matter what.

I can see a possible workaround if custom functions on JMESPath. E.g. if I implement the "if" function for use, I can do something like:

    search("if(bar==`1`, `hello`, `world`)", {'bar': '1'})

This would return the literal `hello` if the `bar` key is 1, otherwise it returns `world`. The only issue is the current python implementation means its going to be hacky to do this. You have to override multiple classes, in `functions.py` and override the `TreeInterpreter` and the `ParsedResult` classes as well. 

I think if custom functions were desired, it would be much more elegant if there is a method to register them directly into the `RuntimeFunctions` in `functions.py`, rather than either forcing a fork or overriding a litany of classes.

What do you think?




Discussion:

Patch:
diff --git a/README.rst b/README.rst
index e7262d4..486783b 100644
--- a/README.rst
+++ b/README.rst
@@ -96,6 +96,85 @@ of your dict keys.  To do this you can use either of these options:
     ...               jmespath.Options(dict_cls=collections.OrderedDict))
 
 
+Custom Functions
+~~~~~~~~~~~~~~~~
+
+The JMESPath language has numerous
+`built-in functions
+<http://jmespath.org/specification.html#built-in-functions>`__, but it is
+also possible to add your own custom functions.  Keep in mind that
+custom function support in jmespath.py is experimental and the API may
+change based on feedback.
+
+**If you have a custom function that you've found useful, consider submitting
+it to jmespath.site and propose that it be added to the JMESPath language.**
+You can submit proposals
+`here <https://github.com/jmespath/jmespath.site/issues>`__.
+
+To create custom functions:
+
+* Create a subclass of ``jmespath.functions.Functions``.
+* Create a method with the name ``_func_<your function name>``.
+* Apply the ``jmespath.functions.signature`` decorator that indicates
+  the expected types of the function arguments.
+* Provide an instance of your subclass in a ``jmespath.Options`` object.
+
+Below are a few examples:
+
+.. code:: python
+
+    import jmespath
+    from jmespath import functions
+
+    # 1. Create a subclass of functions.Functions.
+    #    The function.Functions base class has logic
+    #    that introspects all of its methods and automatically
+    #    registers your custom functions in its function table.
+    class CustomFunctions(functions.Functions):
+
+        # 2 and 3.  Create a function that starts with _func_
+        # and decorate it with @signature which indicates its
+        # expected types.
+        # In this example, we're creating a jmespath function
+        # called "unique_letters" that accepts a single argument
+        # with an expected type "string".
+        @functions.signature({'types': ['string']})
+        def _func_unique_letters(self, s):
+            # Given a string s, return a sorted
+            # string of unique letters: 'ccbbadd' ->  'abcd'
+            return ''.join(sorted(set(s)))
+
+        # Here's another example.  This is creating
+        # a jmespath function called "my_add" that expects
+        # two arguments, both of which should be of type number.
+        @functions.signature({'types': ['number']}, {'types': ['number']})
+        def _func_my_add(self, x, y):
+            return x + y
+
+    # 4. Provide an instance of your subclass in a Options object.
+    options = jmespath.Options(custom_functions=CustomFunctions())
+
+    # Provide this value to jmespath.search:
+    # This will print 3
+    print(
+        jmespath.search(
+            'my_add(`1`, `2`)', {}, options=options)
+    )
+
+    # This will print "abcd"
+    print(
+        jmespath.search(
+            'foo.bar | unique_letters(@)',
+            {'foo': {'bar': 'ccbbadd'}},
+            options=options)
+    )
+
+Again, if you come up with useful functions that you think make
+sense in the JMESPath language (and make sense to implement in all
+JMESPath libraries, not just python), please let us know at
+`jmespath.site <https://github.com/jmespath/jmespath.site/issues>`__.
+
+
 Specification
 =============
 
diff --git a/jmespath/compat.py b/jmespath/compat.py
index 7b70adb..2ed0fe7 100644
--- a/jmespath/compat.py
+++ b/jmespath/compat.py
@@ -3,6 +3,15 @@ import inspect
 
 PY2 = sys.version_info[0] == 2
 
+
+def with_metaclass(meta, *bases):
+    # Taken from flask/six.
+    class metaclass(meta):
+        def __new__(cls, name, this_bases, d):
+            return meta(name, bases, d)
+    return type.__new__(metaclass, 'temporary_class', (), {})
+
+
 if PY2:
     text_type = unicode
     string_type = basestring
diff --git a/jmespath/functions.py b/jmespath/functions.py
index e306f7b..a7c7bd9 100644
--- a/jmespath/functions.py
+++ b/jmespath/functions.py
@@ -1,10 +1,9 @@
 import math
 import json
-import weakref
 
 from jmespath import exceptions
 from jmespath.compat import string_type as STRING_TYPE
-from jmespath.compat import get_methods
+from jmespath.compat import get_methods, with_metaclass
 
 
 # python types -> jmespath types
@@ -35,48 +34,39 @@ REVERSE_TYPES_MAP = {
 }
 
 
-def populate_function_table(cls):
-    func_table = cls.FUNCTION_TABLE
-    for name, method in get_methods(cls):
-        signature = getattr(method, 'signature', None)
-        if signature is not None:
-            func_table[name[6:]] = {"function": method,
-                                    "signature": signature}
-    return cls
-
-
-def builtin_function(*arguments):
-    def _record_arity(func):
+def signature(*arguments):
+    def _record_signature(func):
         func.signature = arguments
         return func
-    return _record_arity
+    return _record_signature
 
 
-@populate_function_table
-class RuntimeFunctions(object):
-    # The built in functions are automatically populated in the FUNCTION_TABLE
-    # using the @builtin_function decorator on methods defined in this class.
+class FunctionRegistry(type):
+    def __init__(cls, name, bases, attrs):
+        cls._populate_function_table()
+        super(FunctionRegistry, cls).__init__(name, bases, attrs)
 
-    FUNCTION_TABLE = {
-    }
+    def _populate_function_table(cls):
+        function_table = getattr(cls, 'FUNCTION_TABLE', {})
+        # Any method with a @signature decorator that also
+        # starts with "_func_" is registered as a function.
+        # _func_max_by -> max_by function.
+        for name, method in get_methods(cls):
+            if not name.startswith('_func_'):
+                continue
+            signature = getattr(method, 'signature', None)
+            if signature is not None:
+                function_table[name[6:]] = {
+                    'function': method,
+                    'signature': signature,
+                }
+        cls.FUNCTION_TABLE = function_table
 
-    def __init__(self):
-        self._interpreter = None
 
-    @property
-    def interpreter(self):
-        if self._interpreter is None:
-            return None
-        else:
-            return self._interpreter()
+class Functions(with_metaclass(FunctionRegistry, object)):
 
-    @interpreter.setter
-    def interpreter(self, value):
-        # A weakref is used because we have
-        # a cyclic reference and we want to allow
-        # for the memory to be properly freed when
-        # the objects are no longer needed.
-        self._interpreter = weakref.ref(value)
+    FUNCTION_TABLE = {
+    }
 
     def call_function(self, function_name, resolved_args):
         try:
@@ -170,28 +160,28 @@ class RuntimeFunctions(object):
                     raise exceptions.JMESPathTypeError(
                         function_name, element, actual_typename, types)
 
-    @builtin_function({'types': ['number']})
+    @signature({'types': ['number']})
     def _func_abs(self, arg):
         return abs(arg)
 
-    @builtin_function({'types': ['array-number']})
+    @signature({'types': ['array-number']})
     def _func_avg(self, arg):
         return sum(arg) / float(len(arg))
 
-    @builtin_function({'types': [], 'variadic': True})
+    @signature({'types': [], 'variadic': True})
     def _func_not_null(self, *arguments):
         for argument in arguments:
             if argument is not None:
                 return argument
 
-    @builtin_function({'types': []})
+    @signature({'types': []})
     def _func_to_array(self, arg):
         if isinstance(arg, list):
             return arg
         else:
             return [arg]
 
-    @builtin_function({'types': []})
+    @signature({'types': []})
     def _func_to_string(self, arg):
         if isinstance(arg, STRING_TYPE):
             return arg
@@ -199,7 +189,7 @@ class RuntimeFunctions(object):
             return json.dumps(arg, separators=(',', ':'),
                               default=str)
 
-    @builtin_function({'types': []})
+    @signature({'types': []})
     def _func_to_number(self, arg):
         if isinstance(arg, (list, dict, bool)):
             return None
@@ -216,88 +206,88 @@ class RuntimeFunctions(object):
             except ValueError:
                 return None
 
-    @builtin_function({'types': ['array', 'string']}, {'types': []})
+    @signature({'types': ['array', 'string']}, {'types': []})
     def _func_contains(self, subject, search):
         return search in subject
 
-    @builtin_function({'types': ['string', 'array', 'object']})
+    @signature({'types': ['string', 'array', 'object']})
     def _func_length(self, arg):
         return len(arg)
 
-    @builtin_function({'types': ['string']}, {'types': ['string']})
+    @signature({'types': ['string']}, {'types': ['string']})
     def _func_ends_with(self, search, suffix):
         return search.endswith(suffix)
 
-    @builtin_function({'types': ['string']}, {'types': ['string']})
+    @signature({'types': ['string']}, {'types': ['string']})
     def _func_starts_with(self, search, suffix):
         return search.startswith(suffix)
 
-    @builtin_function({'types': ['array', 'string']})
+    @signature({'types': ['array', 'string']})
     def _func_reverse(self, arg):
         if isinstance(arg, STRING_TYPE):
             return arg[::-1]
         else:
             return list(reversed(arg))
 
-    @builtin_function({"types": ['number']})
+    @signature({"types": ['number']})
     def _func_ceil(self, arg):
         return math.ceil(arg)
 
-    @builtin_function({"types": ['number']})
+    @signature({"types": ['number']})
     def _func_floor(self, arg):
         return math.floor(arg)
 
-    @builtin_function({"types": ['string']}, {"types": ['array-string']})
+    @signature({"types": ['string']}, {"types": ['array-string']})
     def _func_join(self, separator, array):
         return separator.join(array)
 
-    @builtin_function({'types': ['expref']}, {'types': ['array']})
+    @signature({'types': ['expref']}, {'types': ['array']})
     def _func_map(self, expref, arg):
         result = []
         for element in arg:
-            result.append(self.interpreter.visit(expref.expression, element))
+            result.append(expref.visit(expref.expression, element))
         return result
 
-    @builtin_function({"types": ['array-number', 'array-string']})
+    @signature({"types": ['array-number', 'array-string']})
     def _func_max(self, arg):
         if arg:
             return max(arg)
         else:
             return None
 
-    @builtin_function({"types": ["object"], "variadic": True})
+    @signature({"types": ["object"], "variadic": True})
     def _func_merge(self, *arguments):
         merged = {}
         for arg in arguments:
             merged.update(arg)
         return merged
 
-    @builtin_function({"types": ['array-number', 'array-string']})
+    @signature({"types": ['array-number', 'array-string']})
     def _func_min(self, arg):
         if arg:
             return min(arg)
         else:
             return None
 
-    @builtin_function({"types": ['array-string', 'array-number']})
+    @signature({"types": ['array-string', 'array-number']})
     def _func_sort(self, arg):
         return list(sorted(arg))
 
-    @builtin_function({"types": ['array-number']})
+    @signature({"types": ['array-number']})
     def _func_sum(self, arg):
         return sum(arg)
 
-    @builtin_function({"types": ['object']})
+    @signature({"types": ['object']})
     def _func_keys(self, arg):
         # To be consistent with .values()
         # should we also return the indices of a list?
         return list(arg.keys())
 
-    @builtin_function({"types": ['object']})
+    @signature({"types": ['object']})
     def _func_values(self, arg):
         return list(arg.values())
 
-    @builtin_function({'types': []})
+    @signature({'types': []})
     def _func_type(self, arg):
         if isinstance(arg, STRING_TYPE):
             return "string"
@@ -312,7 +302,7 @@ class RuntimeFunctions(object):
         elif arg is None:
             return "null"
 
-    @builtin_function({'types': ['array']}, {'types': ['expref']})
+    @signature({'types': ['array']}, {'types': ['expref']})
     def _func_sort_by(self, array, expref):
         if not array:
             return array
@@ -323,34 +313,32 @@ class RuntimeFunctions(object):
         # that validates that type, which requires that remaining array
         # elements resolve to the same type as the first element.
         required_type = self._convert_to_jmespath_type(
-            type(self.interpreter.visit(expref.expression, array[0])).__name__)
+            type(expref.visit(expref.expression, array[0])).__name__)
         if required_type not in ['number', 'string']:
             raise exceptions.JMESPathTypeError(
                 'sort_by', array[0], required_type, ['string', 'number'])
-        keyfunc = self._create_key_func(expref.expression,
+        keyfunc = self._create_key_func(expref,
                                         [required_type],
                                         'sort_by')
         return list(sorted(array, key=keyfunc))
 
-    @builtin_function({'types': ['array']}, {'types': ['expref']})
+    @signature({'types': ['array']}, {'types': ['expref']})
     def _func_min_by(self, array, expref):
-        keyfunc = self._create_key_func(expref.expression,
+        keyfunc = self._create_key_func(expref,
                                         ['number', 'string'],
                                         'min_by')
         return min(array, key=keyfunc)
 
-    @builtin_function({'types': ['array']}, {'types': ['expref']})
+    @signature({'types': ['array']}, {'types': ['expref']})
     def _func_max_by(self, array, expref):
-        keyfunc = self._create_key_func(expref.expression,
+        keyfunc = self._create_key_func(expref,
                                         ['number', 'string'],
                                         'min_by')
         return max(array, key=keyfunc)
 
-    def _create_key_func(self, expr_node, allowed_types, function_name):
-        interpreter = self.interpreter
-
+    def _create_key_func(self, expref, allowed_types, function_name):
         def keyfunc(x):
-            result = interpreter.visit(expr_node, x)
+            result = expref.visit(expref.expression, x)
             actual_typename = type(result).__name__
             jmespath_type = self._convert_to_jmespath_type(actual_typename)
             # allowed_types is in term of jmespath types, not python types.
diff --git a/jmespath/visitor.py b/jmespath/visitor.py
index bbb5a8e..8a94a7e 100644
--- a/jmespath/visitor.py
+++ b/jmespath/visitor.py
@@ -35,7 +35,7 @@ def _is_special_integer_case(x, y):
 
 class Options(object):
     """Options to control how a JMESPath function is evaluated."""
-    def __init__(self, dict_cls):
+    def __init__(self, dict_cls=None, custom_functions=None):
         #: The class to use when creating a dict.  The interpreter
         #  may create dictionaries during the evalution of a JMESPath
         #  expression.  For example, a multi-select hash will
@@ -45,11 +45,16 @@ class Options(object):
         #  want to set a collections.OrderedDict so that you can
         #  have predictible key ordering.
         self.dict_cls = dict_cls
+        self.custom_functions = custom_functions
 
 
 class _Expression(object):
-    def __init__(self, expression):
+    def __init__(self, expression, interpreter):
         self.expression = expression
+        self.interpreter = interpreter
+
+    def visit(self, node, *args, **kwargs):
+        return self.interpreter.visit(node, *args, **kwargs)
 
 
 class Visitor(object):
@@ -83,15 +88,16 @@ class TreeInterpreter(Visitor):
 
     def __init__(self, options=None):
         super(TreeInterpreter, self).__init__()
-        self._options = options
         self._dict_cls = self.MAP_TYPE
-        if options is not None and options.dict_cls is not None:
+        if options is None:
+            options = Options()
+        self._options = options
+        if options.dict_cls is not None:
             self._dict_cls = self._options.dict_cls
-        self._functions = functions.RuntimeFunctions()
-        # Note that .interpreter is a property that uses
-        # a weakref so that the cyclic reference can be
-        # properly freed.
-        self._functions.interpreter = self
+        if options.custom_functions is not None:
+            self._functions = self._options.custom_functions
+        else:
+            self._functions = functions.Functions()
 
     def default_visit(self, node, *args, **kwargs):
         raise NotImplementedError(node['type'])
@@ -119,7 +125,7 @@ class TreeInterpreter(Visitor):
         return value
 
     def visit_expref(self, node, value):
-        return _Expression(node['children'][0])
+        return _Expression(node['children'][0], self)
 
     def visit_function_expression(self, node, value):
         resolved_args = []


Test patch:
diff --git a/tests/test_search.py b/tests/test_search.py
index 56a0a75..71ab3dc 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -1,6 +1,7 @@
 from tests import unittest, OrderedDict
 
 import jmespath
+import jmespath.functions
 
 
 class TestSearchOptions(unittest.TestCase):
@@ -10,3 +11,28 @@ class TestSearchOptions(unittest.TestCase):
             {'c': 'c', 'b': 'b', 'a': 'a', 'd': 'd'},
             options=jmespath.Options(dict_cls=OrderedDict))
         self.assertEqual(result, ['a', 'b', 'c'])
+
+    def test_can_provide_custom_functions(self):
+        class CustomFunctions(jmespath.functions.Functions):
+            @jmespath.functions.signature(
+                {'types': ['number']},
+                {'types': ['number']})
+            def _func_custom_add(self, x, y):
+                return x + y
+
+            @jmespath.functions.signature(
+                {'types': ['number']},
+                {'types': ['number']})
+            def _func_my_subtract(self, x, y):
+                return x - y
+
+
+        options = jmespath.Options(custom_functions=CustomFunctions())
+        self.assertEqual(
+            jmespath.search('custom_add(`1`, `2`)', {}, options=options),
+            3
+        )
+        self.assertEqual(
+            jmespath.search('my_subtract(`10`, `3`)', {}, options=options),
+            7
+        )