Add support for custom functions

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

jmespath/compat.py

@@ -3,6 +3,15 @@
 
 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

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 @@
 }
 
 
-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,36 +160,36 @@ def _subtype_check(self, current, allowed_subtypes, types, function_name):
                     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
         else:
             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 @@ def _func_to_number(self, arg):
             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 @@ def _func_type(self, arg):
         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 @@ def _func_sort_by(self, array, expref):
         # 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.