Add search_json(): query a JSON document given as a string, with optional native accelerator

CHANGELOG.rst

@@ -1,7 +1,10 @@
 Next Release (TBD)
 ==================
 
-* No changes yet.
+* Added ``jmespath.search_json()`` to query a JSON document given as a
+  string. When the optional ``aero-jmespath`` native accelerator is installed
+  the whole pipeline runs in native code; otherwise it falls back to
+  ``json.loads()`` + :func:`search`.
 
 
 1.1.0

jmespath/__init__.py

@@ -1,12 +1,51 @@
+import json
+
 from jmespath import parser
+from jmespath.compat import string_type
 from jmespath.visitor import Options
 
 __version__ = '1.1.0'
 
+# Optional native accelerator (e.g. ``aero-jmespath``). When installed it exposes
+# ``search(expression, json_string)`` that runs the whole parse + eval + serialize
+# pipeline in native code.  ``search_json`` uses it automatically when ``data`` is
+# a JSON string and falls back to pure Python otherwise, so this is purely an
+# opt-in performance improvement and never changes behaviour.
+_NATIVE_ACCELERATOR = None
+try:
+    import aero_jmespath as _accelerator
+    _NATIVE_ACCELERATOR = _accelerator.search
+except ImportError:
+    pass
+
 
 def compile(expression):
     return parser.Parser().parse(expression)
 
 
 def search(expression, data, options=None):
     return parser.Parser().parse(expression).search(data, options=options)
+
+
+def search_json(expression, data, options=None):
+    """Search a JSON document given as a *string*.
+
+    ``data`` must be a JSON-encoded string.  When the optional native
+    accelerator (``aero-jmespath``) is installed, the expression is evaluated
+    entirely in native code and the serialized result is decoded back to a
+    Python value.  Otherwise this is equivalent to
+    ``search(expression, json.loads(data))``.
+
+    The return value is the same as :func:`search`.
+    """
+    if _NATIVE_ACCELERATOR is not None and isinstance(data, string_type):
+        result = _NATIVE_ACCELERATOR(expression, data)
+        # The native kernel reports errors (bad JSON, unsupported expression,
+        # invalid value) as "\x1eERR<code>". Falling back to pure Python here
+        # keeps exceptions identical to the non-accelerated path (ValueError
+        # for bad JSON, ParseError for a bad expression).
+        if not result.startswith(b"\x1eERR"):
+            return json.loads(result)
+    if isinstance(data, string_type):
+        data = json.loads(data)
+    return search(expression, data, options=options)

tests/test_search.py

@@ -1,5 +1,6 @@
 import sys
 import decimal
+import json
 from tests import unittest, OrderedDict
 
 import jmespath
@@ -62,3 +63,39 @@ def test_can_handle_decimals_as_numeric_type(self):
         result = decimal.Decimal('3')
         self.assertEqual(jmespath.search('[?a >= `1`].a', [{'a': result}]),
                          [result])
+
+
+class TestSearchJson(unittest.TestCase):
+    """search_json() takes a JSON string and returns the same value as search()."""
+
+    def test_basic_field(self):
+        self.assertEqual(
+            jmespath.search_json('a.b', '{"a": {"b": "x"}}'),
+            'x')
+
+    def test_projection(self):
+        self.assertEqual(
+            jmespath.search_json('a[*].b', '{"a": [{"b": 1}, {"b": 2}]}'),
+            [1, 2])
+
+    def test_empty_doc(self):
+        self.assertEqual(
+            jmespath.search_json('a', '{}'),
+            None)
+
+    def test_invalid_json_raises(self):
+        with self.assertRaises(ValueError):
+            jmespath.search_json('a', '{not json}')
+
+    def test_matches_search_on_parsed_doc(self):
+        doc = '{"servers": [{"name": "x", "up": true}, {"name": "y", "up": false}]}'
+        expr = 'servers[?up == `true`].name'
+        self.assertEqual(
+            jmespath.search_json(expr, doc),
+            jmespath.search(expr, json.loads(doc)))
+
+    def test_passes_options_through(self):
+        self.assertEqual(
+            jmespath.search_json('a.b', '{"a": {"b": "x"}}',
+                                 options=jmespath.Options()),
+            'x')