Merge pull request #23 from linw1995/hotfix/slice_nested_exprs_find_e…

…rror

Hotfix/slice nested exprs find error

jsonpath/core.py

@@ -570,9 +570,9 @@ class Slice(Expr):
 
     def __init__(
         self,
-        start: Optional[Union[Expr, int]] = None,
-        stop: Optional[Union[Expr, int]] = None,
-        step: Optional[Union[Expr, int]] = None,
+        start: Union[Expr, int, None] = None,
+        stop: Union[Expr, int, None] = None,
+        step: Union[Expr, int, None] = None,
     ) -> None:
         super().__init__()
         self.start = start
@@ -596,49 +596,34 @@ def _get_partial_expression(self) -> str:
 
         return ":".join(parts)
 
-    def find(self, element: List[Any]) -> Any:
-        if isinstance(element, list):
-            # set var_finding False to start new finding process for
-            # the nested expr: self.start, self.end and self.step
+    def _ensure_int_or_none(
+        self, value: Union[Expr, int, None]
+    ) -> Union[int, None]:
+        if isinstance(value, Expr):
+            # set var_finding False to start new finding process for the nested expr
             with temporary_set(var_finding, False):
-                start = (
-                    self.start.find(element)
-                    if isinstance(self.start, Expr)
-                    else self.start
-                )
-                end = (
-                    self.end.find(element)
-                    if isinstance(self.end, Expr)
-                    else self.end
-                )
-                step = (
-                    self.step.find(element)
-                    if isinstance(self.step, Expr)
-                    else self.step
-                )
-
-            if not start:
-                start = 0
-            elif isinstance(start, list):
-                start = start[0]
-                if not isinstance(start, int):
-                    return []
-            if not end:
-                end = len(element)
-            elif isinstance(end, list):
-                end = end[0]
-                if not isinstance(end, int):
-                    return []
-            if not step:
-                step = 1
-            elif isinstance(step, list):
-                step = step[0]
-                if not isinstance(step, int):
-                    return []
+                found_elements = value.find(var_parent.get())
+            if not found_elements or not isinstance(found_elements[0], int):
+                raise JSONPathFindError
+            return found_elements[0]
+        else:
+            return value
 
-            return element[start:end:step]
+    def find(self, element: List[Any]) -> List[Any]:
+        assert isinstance(element, list), "Slice.find apply on list only."
 
-        raise JSONPathFindError
+        start = self._ensure_int_or_none(self.start)
+        end = self._ensure_int_or_none(self.end)
+        step = self._ensure_int_or_none(self.step)
+
+        if start is None:
+            start = 0
+        if end is None:
+            end = len(element)
+        if step is None:
+            step = 1
+
+        return element[start:end:step]
 
 
 class Brace(Expr):

tests/test_lark.py

@@ -143,6 +143,13 @@ def test_parse_check_and_extract(expression, data, expect):
         ("$[:]", [1, 2], [1, 2]),
         ("$[:-1]", [1, 2, 3], [1, 2]),
         ("$[::2]", [1, 2, 3], [1, 3]),
+        ("$[a:]", [1, 2], []),
+        ("$[:a]", [1, 2], []),
+        ("$[::a]", [1, 2], []),
+        ("$.data[$.a:]", {"data": [1, 2]}, []),
+        ("$.data[$.a:]", {"data": [1, 2], "a": 1}, [2]),
+        ("$.data[a:]", {"data": [1, 2], "a": 1}, [2]),
+        ("$.data[a:b]", {"data": [1, 2], "a": 1, "b": 1}, []),
         ("$.*[0]", {"boo": [1, 2, 3], "bar": [2, 3, 4]}, [1, 2]),
         ("$.*[*]", {"boo": [1, 2, 3], "bar": [2, 3, 4]}, [1, 2, 3, 2, 3, 4]),
         ("($.*)[0]", {"boo": [1, 2, 3], "bar": [2, 3, 4]}, [[1, 2, 3]]),