Add support for 'arraylike' objects as JSON arrays

CHANGELOG.rst

@@ -1,3 +1,8 @@
+Unreleased
+==========
+
+* Added support for Arraylike python objects as json arrays.
+
 1.0.1
 =====
 

jmespath/functions.py

@@ -1,5 +1,6 @@
 import math
 import json
+from typing import Sequence
 
 from jmespath import exceptions
 from jmespath.compat import string_type as STRING_TYPE
@@ -35,6 +36,14 @@
 }
 
 
+def is_array(arg):
+ return hasattr(arg, "__array__") and arg.shape != ()
+
+
+def is_arraylike(arg):
+ return (isinstance(arg, Sequence) and not isinstance(arg, (str, bytes))) or is_array(arg)
+
+
 def signature(*arguments):
  def _record_signature(func):
  func.signature = arguments
@@ -180,7 +189,7 @@ def _func_not_null(self, *arguments):
 
  @signature({'types': []})
  def _func_to_array(self, arg):
- if isinstance(arg, list):
+ if is_arraylike(arg):
  return arg
  else:
  return [arg]
@@ -297,7 +306,7 @@ def _func_type(self, arg):
  return "string"
  elif isinstance(arg, bool):
  return "boolean"
- elif isinstance(arg, list):
+ elif is_arraylike(arg):
  return "array"
  elif isinstance(arg, dict):
  return "object"

jmespath/visitor.py

@@ -1,15 +1,26 @@
 import operator
 
 from jmespath import functions
+from jmespath.functions import is_array, is_arraylike
 from jmespath.compat import string_type
 from numbers import Number
 
 
+def _arraylike_all(arg):
+ return arg.__array__().all() if is_array(arg) else arg
+
+
+def _arraylike_to_list(arg):
+ return [_arraylike_to_list(i) for i in arg] if is_arraylike(arg) else arg
+
+
 def _equals(x, y):
  if _is_special_number_case(x, y):
  return False
+ elif is_array(x) or is_array(y):
+ return _arraylike_all(x == y)
  else:
- return x == y
+ return _arraylike_to_list(x) == _arraylike_to_list(y)
 
 
 def _is_special_number_case(x, y):
@@ -172,7 +183,7 @@ def visit_function_expression(self, node, value):
 
  def visit_filter_projection(self, node, value):
  base = self.visit(node['children'][0], value)
- if not isinstance(base, list):
+ if not is_arraylike(base):
  return None
  comparator_node = node['children'][2]
  collected = []
@@ -185,12 +196,12 @@ def visit_filter_projection(self, node, value):
 
  def visit_flatten(self, node, value):
  base = self.visit(node['children'][0], value)
- if not isinstance(base, list):
- # Can't flatten the object if it's not a list.
+ if not is_arraylike(base):
+ # Can't flatten the object if it's not arraylike.
  return None
  merged_list = []
  for element in base:
- if isinstance(element, list):
+ if is_arraylike(element):
  merged_list.extend(element)
  else:
  merged_list.append(element)
@@ -202,7 +213,7 @@ def visit_identity(self, node, value):
  def visit_index(self, node, value):
  # Even though we can index strings, we don't
  # want to support that.
- if not isinstance(value, list):
+ if not is_arraylike(value):
  return None
  try:
  return value[node['value']]
@@ -216,7 +227,7 @@ def visit_index_expression(self, node, value):
  return result
 
  def visit_slice(self, node, value):
- if not isinstance(value, list):
+ if not is_arraylike(value):
  return None
  s = slice(*node['children'])
  return value[s]
@@ -271,7 +282,7 @@ def visit_pipe(self, node, value):
 
  def visit_projection(self, node, value):
  base = self.visit(node['children'][0], value)
- if not isinstance(base, list):
+ if not is_arraylike(base):
  return None
  collected = []
  for element in base:

requirements.txt

@@ -1,6 +1,11 @@
 wheel==0.38.1
+parameterized==0.9.0
 pytest==6.2.5
 pytest-cov==3.0.0
 hypothesis==3.1.0 ; python_version < '3.8'
 hypothesis==5.5.4 ; python_version == '3.8'
 hypothesis==5.35.4 ; python_version == '3.9'
+astropy>=3.1
+dask>=2.0.0
+numpy>=1.15.0
+xarray>=0.18.0

tests/test_arraylike.py

@@ -0,0 +1,132 @@
+import astropy.units as u
+import dask.array as da
+import numpy as np
+import xarray as xr
+from parameterized import parameterized, parameterized_class
+
+import jmespath
+import jmespath.functions
+from tests import unittest
+
+
+@parameterized_class(("name", "data"), [
+ ("list", {
+ "value": {
+ "data": [[1,2,3],[4,5,6],[7,8,9]]
+ },
+ "same": {
+ "data": [[1,2,3],[4,5,6],[7,8,9]]
+ },
+ "other": {
+ "data": [[2,2,3],[4,5,6],[7,8,9]]
+ }
+ }),
+ ("tuple", {
+ "value": {
+ "data": ((1,2,3),(4,5,6),(7,8,9))
+ },
+ "same": {
+ "data": ([1,2,3],[4,5,6],[7,8,9])
+ },
+ "other": {
+ "data": [[2,2,3],[4,5,6],[7,8,9]]
+ }
+ }),
+ ("numpy", {
+ "value": {
+ "data": np.array([[1,2,3],[4,5,6],[7,8,9]])
+ },
+ "same": {
+ "data": (np.array([1,2,3]),np.array([4,5,6]),np.array([7,8,9]))
+ },
+ "other": {
+ "data": np.array([[2,2,3],[4,5,6],[7,8,9]])
+ }
+ }),
+ ("dask", {
+ "value": {
+ "data": da.from_array([[1,2,3],[4,5,6],[7,8,9]])
+ },
+ "same": {
+ "data": (da.from_array([1,2,3]),da.from_array([4,5,6]),da.from_array([7,8,9]))
+ },
+ "other": {
+ "data": da.from_array([[2,2,3],[4,5,6],[7,8,9]])
+ }
+ }),
+ ("xarray", {
+ "value": {
+ "data": xr.DataArray([[1,2,3],[4,5,6],[7,8,9]])
+ },
+ "same": {
+ "data": (xr.DataArray([1,2,3]),xr.DataArray([4,5,6]),xr.DataArray([7,8,9]))
+ },
+ "other": {
+ "data": xr.DataArray([[2,2,3],[4,5,6],[7,8,9]])
+ }
+ }),
+ ("astropy", {
+ "value": {
+ "data": u.Quantity([[1,2,3],[4,5,6],[7,8,9]])
+ },
+ "same": {
+ "data": (u.Quantity([1,2,3]),u.Quantity([4,5,6]),u.Quantity([7,8,9]))
+ },
+ "other": {
+ "data": u.Quantity([[2,2,3],[4,5,6],[7,8,9]])
+ }
+ }),
+])
+class TestArrayNumeric(unittest.TestCase):
+ @parameterized.expand([
+ ["self", "@", lambda data: data],
+ ["get", "value.data", lambda data: data["value"]["data"]],
+ ["slice_horizontal", "value.data[1][:]", lambda data: np.array(data["value"]["data"])[1,:]],
+ ["slice_horizontal2", "value.data[:3:2][:]", lambda data: np.array(data["value"]["data"])[:3:2,:]],
+ ["slice_vertical", "value.data[:][1]", lambda data: np.array(data["value"]["data"])[:,1]],
+ ["slice_vertical2", "value.data[:][:3:2]", lambda data: np.array(data["value"]["data"])[:,:3:2]],
+ ["flatten", "value.data[]", lambda data: np.array(data["value"]["data"]).flatten()],
+ ["compare_self", "value.data == value.data", lambda _: True],
+ ["compare_same", "value.data == same.data", lambda _: True],
+ ["compare_other", "value.data == other.data", lambda _: False],
+ ["compare_literal_scalar", "value.data[0][0] == `1`", lambda _: True],
+ ["compare_literal_slice", "value.data[1][:] == `[4, 5, 6]`", lambda _: True],
+ ["compare_literal", "value.data == `[[1,2,3],[4,5,6],[7,8,9]]`", lambda _: True],
+ ["compare_flattened", "value.data[] == `[1,2,3,4,5,6,7,8,9]`", lambda _: True],
+ ])
+ def test_search(self, test_name, query, expected):
+ result = jmespath.search(query, self.data)
+ np.testing.assert_array_equal(result, expected(self.data), test_name)
+
+
+@parameterized_class(("name", "data"), [
+ ("numpy", {
+ "value": { 
+ "data": np.array([["test", "messages"],["in", "numpy"]])
+ },
+ "same": { 
+ "data": np.array([["test", "messages"],["in", "numpy"]])
+ },
+ "other": { 
+ "data": np.array([["test", "messages"],["other", "numpy"]])
+ }
+ })
+])
+class TestArrayStr(unittest.TestCase):
+ @parameterized.expand([
+ ["self", "@", lambda data: data],
+ ["get", "value.data", lambda data: data["value"]["data"]],
+ ["slice_horizontal", "value.data[1][:]", lambda data: data["value"]["data"][1,:]],
+ ["slice_vertical", "value.data[:][1]", lambda data: data["value"]["data"][:,1]],
+ ["flatten", "value.data[]", lambda data: data["value"]["data"].flatten()],
+ ["compare_self", "value.data == value.data", lambda _: True],
+ ["compare_same", "value.data == same.data", lambda _: True],
+ ["compare_other", "value.data == other.data", lambda _: False],
+ ["compare_literal_scalar", "value.data[0][0] == 'test'", lambda _: True],
+ ["compare_literal_slice", "value.data[1][:] == ['in', 'numpy']", lambda _: True],
+ ["compare_literal", "value.data == [['test', 'messages'],['in', 'numpy']]", lambda _: True],
+ ["compare_flattened", "value.data[] == ['test', 'messages', 'in', 'numpy']", lambda _: True],
+ ])
+ def test_search(self, name, query, expected):
+ result = jmespath.search(query, self.data)
+ np.testing.assert_array_equal(result, expected(self.data), name)