Add support for custom functions

Closes #100.

Author: jamesls

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

@@ -3,7 +3,7 @@
 
 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
@@ -34,27 +34,36 @@
 }
 
 
-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 signature(*arguments):
     def _record_signature(func):
         func.signature = arguments
         return func
     return _record_signature
 
 
-@populate_function_table
-class RuntimeFunctions(object):
-    # The built in functions are automatically populated in the FUNCTION_TABLE
-    # using the @signature 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)
+
+    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
+
+
+class Functions(with_metaclass(FunctionRegistry, object)):
 
     FUNCTION_TABLE = {
     }

jmespath/visitor.py

@@ -97,7 +97,7 @@ def __init__(self, options=None):
         if options.custom_functions is not None:
             self._functions = self._options.custom_functions
         else:
-            self._functions = functions.RuntimeFunctions()
+            self._functions = functions.Functions()
 
     def default_visit(self, node, *args, **kwargs):
         raise NotImplementedError(node['type'])

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 @@ def test_can_provide_dict_cls(self):
             {'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
+        )