Add "key" and "keys filter" JSONPath selectors

Author: jg-rp

jsonpath/env.py

@@ -102,6 +102,8 @@ class attributes `root_token`, `self_token` and `filter_context_token`.
             filtering a mapping or sequence. Defaults to `"#"`.
         keys_selector_token (str): The pattern used as the "keys" selector. Defaults to
             `"~"`.
+        keys_filter_token (str): The pattern used as the "keys filter" selector.
+            Defaults to `"~?"`.
         lexer_class: The lexer to use when tokenizing path strings.
         max_int_index (int): The maximum integer allowed when selecting array items by
             index. Defaults to `(2**53) - 1`.
@@ -122,6 +124,7 @@ class attributes `root_token`, `self_token` and `filter_context_token`.
     intersection_token = "&"
     key_token = "#"
     keys_selector_token = "~"
+    keys_filter_token = "~?"
     root_token = "$"
     self_token = "@"
     union_token = "|"

jsonpath/lex.py

@@ -15,6 +15,7 @@
 from .token import TOKEN_CONTAINS
 from .token import TOKEN_DDOT
 from .token import TOKEN_DOT
+from .token import TOKEN_DOT_KEY_PROPERTY
 from .token import TOKEN_DOT_PROPERTY
 from .token import TOKEN_DOUBLE_QUOTE_STRING
 from .token import TOKEN_EQ
@@ -30,6 +31,7 @@
 from .token import TOKEN_INT
 from .token import TOKEN_INTERSECTION
 from .token import TOKEN_KEY
+from .token import TOKEN_KEY_NAME
 from .token import TOKEN_KEYS
 from .token import TOKEN_KEYS_FILTER
 from .token import TOKEN_LBRACKET
@@ -103,6 +105,13 @@ def __init__(self, *, env: JSONPathEnvironment) -> None:
         # .thing
         self.dot_property_pattern = rf"(?P<G_DOT>\.)(?P<G_PROP>{self.key_pattern})"
 
+        # .~thing
+        self.dot_key_pattern = (
+            r"(?P<G_DOT_KEY>\.)"
+            rf"(?P<G_KEY>{re.escape(env.keys_selector_token)})"
+            rf"(?P<G_PROP_KEY>{self.key_pattern})"
+        )
+
         # /pattern/ or /pattern/flags
         self.re_pattern = r"/(?P<G_RE>.+?)/(?P<G_RE_FLAGS>[aims]*)"
 
@@ -122,12 +131,14 @@ def compile_rules(self) -> Pattern[str]:
             (TOKEN_INTERSECTION, self.env.intersection_token),
             (TOKEN_FILTER_CONTEXT, self.env.filter_context_token),
             (TOKEN_KEYS, self.env.keys_selector_token),
+            (TOKEN_KEYS_FILTER, self.env.keys_filter_token),
         ]
 
         rules = [
             (TOKEN_DOUBLE_QUOTE_STRING, self.double_quote_pattern),
             (TOKEN_SINGLE_QUOTE_STRING, self.single_quote_pattern),
             (TOKEN_RE_PATTERN, self.re_pattern),
+            (TOKEN_DOT_KEY_PROPERTY, self.dot_key_pattern),
             (TOKEN_DOT_PROPERTY, self.dot_property_pattern),
             (TOKEN_FLOAT, r"-?\d+\.\d*(?:[eE][+-]?\d+)?"),
             (TOKEN_INT, r"-?\d+(?P<G_EXP>[eE][+\-]?\d+)?\b"),
@@ -144,7 +155,6 @@ def compile_rules(self) -> Pattern[str]:
             ],
             (TOKEN_WILD, r"\*"),
             (TOKEN_FILTER, r"\?"),
-            (TOKEN_KEYS_FILTER, r"~\?"),  # TODO: get from env
             (TOKEN_IN, r"in\b"),
             (TOKEN_TRUE, r"[Tt]rue\b"),
             (TOKEN_FALSE, r"[Ff]alse\b"),
@@ -199,6 +209,17 @@ def tokenize(self, path: str) -> Iterator[Token]:  # noqa PLR0912
                     value=match.group("G_PROP"),
                     index=match.start("G_PROP"),
                 )
+            elif kind == TOKEN_DOT_KEY_PROPERTY:
+                yield _token(
+                    kind=TOKEN_DOT,
+                    value=match.group("G_DOT_KEY"),
+                    index=match.start("G_DOT_KEY"),
+                )
+                yield _token(
+                    kind=TOKEN_KEY_NAME,
+                    value=match.group("G_PROP_KEY"),
+                    index=match.start("G_PROP_KEY"),
+                )
             elif kind == TOKEN_DOUBLE_QUOTE_STRING:
                 yield _token(
                     kind=TOKEN_DOUBLE_QUOTE_STRING,

jsonpath/parse.py

@@ -11,6 +11,7 @@
 from typing import Iterator
 from typing import List
 from typing import Optional
+from typing import Union
 
 from jsonpath.function_extensions.filter_function import ExpressionType
 from jsonpath.function_extensions.filter_function import FilterFunction
@@ -45,6 +46,8 @@
 from .selectors import Filter
 from .selectors import IndexSelector
 from .selectors import JSONPathSelector
+from .selectors import KeySelector
+from .selectors import KeysFilter
 from .selectors import KeysSelector
 from .selectors import PropertySelector
 from .selectors import SliceSelector
@@ -69,7 +72,9 @@
 from .token import TOKEN_INT
 from .token import TOKEN_INTERSECTION
 from .token import TOKEN_KEY
+from .token import TOKEN_KEY_NAME
 from .token import TOKEN_KEYS
+from .token import TOKEN_KEYS_FILTER
 from .token import TOKEN_LBRACKET
 from .token import TOKEN_LE
 from .token import TOKEN_LG
@@ -314,15 +319,15 @@ def parse_path(self, stream: TokenStream) -> Iterable[JSONPathSegment]:
             if _token.kind == TOKEN_DOT:
                 stream.eat(TOKEN_DOT)
                 # Assert that dot is followed by shorthand selector without whitespace.
-                stream.expect(TOKEN_NAME, TOKEN_WILD, TOKEN_KEYS)
+                stream.expect(TOKEN_NAME, TOKEN_WILD, TOKEN_KEYS, TOKEN_KEY_NAME)
                 token = stream.current()
-                selectors = self.parse_selectors(stream)
+                selectors = self.parse_selector(stream)
                 yield JSONPathChildSegment(
                     env=self.env, token=token, selectors=selectors
                 )
             elif _token.kind == TOKEN_DDOT:
                 token = stream.eat(TOKEN_DDOT)
-                selectors = self.parse_selectors(stream)
+                selectors = self.parse_selector(stream)
                 if not selectors:
                     raise JSONPathSyntaxError(
                         "missing selector for recursive descent segment",
@@ -332,22 +337,22 @@ def parse_path(self, stream: TokenStream) -> Iterable[JSONPathSegment]:
                     env=self.env, token=token, selectors=selectors
                 )
             elif _token.kind == TOKEN_LBRACKET:
-                selectors = self.parse_selectors(stream)
+                selectors = self.parse_selector(stream)
                 yield JSONPathChildSegment(
                     env=self.env, token=_token, selectors=selectors
                 )
-            elif _token.kind in {TOKEN_NAME, TOKEN_WILD, TOKEN_KEYS}:
+            elif _token.kind in {TOKEN_NAME, TOKEN_WILD, TOKEN_KEYS, TOKEN_KEY_NAME}:
                 # A non-standard "bare" path. One without a leading identifier (`$`,
                 # `@`, `^` or `_`).
                 token = stream.current()
-                selectors = self.parse_selectors(stream)
+                selectors = self.parse_selector(stream)
                 yield JSONPathChildSegment(
                     env=self.env, token=token, selectors=selectors
                 )
             else:
                 break
 
-    def parse_selectors(self, stream: TokenStream) -> tuple[JSONPathSelector, ...]:
+    def parse_selector(self, stream: TokenStream) -> tuple[JSONPathSelector, ...]:
         token = stream.next()
 
         if token.kind == TOKEN_NAME:
@@ -359,6 +364,15 @@ def parse_selectors(self, stream: TokenStream) -> tuple[JSONPathSelector, ...]:
                 ),
             )
 
+        if token.kind == TOKEN_KEY_NAME:
+            return (
+                KeySelector(
+                    env=self.env,
+                    token=token,
+                    key=token.value,
+                ),
+            )
+
         if token.kind == TOKEN_WILD:
             return (
                 WildSelector(
@@ -432,6 +446,8 @@ def parse_bracketed_selection(self, stream: TokenStream) -> List[JSONPathSelecto
                 stream.next()
             elif token.kind == TOKEN_FILTER:
                 selectors.append(self.parse_filter_selector(stream))
+            elif token.kind == TOKEN_KEYS_FILTER:
+                selectors.append(self.parse_filter_selector(stream, keys=True))
             elif token.kind == TOKEN_EOF:
                 raise JSONPathSyntaxError("unexpected end of query", token=token)
             else:
@@ -514,8 +530,10 @@ def _maybe_index(token: Token) -> bool:
             step=step,
         )
 
-    def parse_filter_selector(self, stream: TokenStream) -> Filter:
-        token = stream.eat(TOKEN_FILTER)
+    def parse_filter_selector(
+        self, stream: TokenStream, *, keys: bool = False
+    ) -> Union[Filter, KeysFilter]:
+        token = stream.next()
         expr = self.parse_filter_expression(stream)
 
         if self.env.well_typed and isinstance(expr, FunctionExtension):
@@ -536,6 +554,11 @@ def parse_filter_selector(self, stream: TokenStream) -> Filter:
                 token=token,
             )
 
+        if keys:
+            return KeysFilter(
+                env=self.env, token=token, expression=BooleanExpression(expr)
+            )
+
         return Filter(env=self.env, token=token, expression=BooleanExpression(expr))
 
     def parse_boolean(self, stream: TokenStream) -> FilterExpression:

jsonpath/selectors.py

@@ -176,10 +176,57 @@ async def resolve_async(self, node: JSONPathMatch) -> AsyncIterable[JSONPathMatc
                 yield match
 
 
+class KeySelector(JSONPathSelector):
+    """Select a single mapping/object name/key.
+
+    NOTE: This is a non-standard selector.
+
+    See https://jg-rp.github.io/json-p3/guides/jsonpath-extra#key-selector.
+    """
+
+    __slots__ = ("key",)
+
+    def __init__(self, *, env: JSONPathEnvironment, token: Token, key: str) -> None:
+        super().__init__(env=env, token=token)
+        self.key = key
+
+    def __str__(self) -> str:
+        return f"{self.env.keys_selector_token}{canonical_string(self.key)}"
+
+    def __eq__(self, __value: object) -> bool:
+        return (
+            isinstance(__value, KeySelector)
+            and self.token == __value.token
+            and self.key == __value.key
+        )
+
+    def __hash__(self) -> int:
+        return hash((self.token, self.key))
+
+    def resolve(self, node: JSONPathMatch) -> Iterable[JSONPathMatch]:
+        if isinstance(node.obj, Mapping) and self.key in node.obj:
+            match = node.__class__(
+                filter_context=node.filter_context(),
+                obj=self.key,
+                parent=node,
+                parts=node.parts + (f"{self.env.keys_selector_token}{self.key}",),
+                path=f"{node.path}[{self}]",
+                root=node.root,
+            )
+            node.add_child(match)
+            yield match
+
+    async def resolve_async(self, node: JSONPathMatch) -> AsyncIterable[JSONPathMatch]:
+        for _node in self.resolve(node):
+            yield _node
+
+
 class KeysSelector(JSONPathSelector):
     """Select mapping/object keys/properties.
 
     NOTE: This is a non-standard selector.
+
+    See https://jg-rp.github.io/json-p3/guides/jsonpath-extra#keys-selector
     """
 
     __slots__ = ()
@@ -198,13 +245,13 @@ def __hash__(self) -> int:
 
     def _keys(self, node: JSONPathMatch) -> Iterable[JSONPathMatch]:
         if isinstance(node.obj, Mapping):
-            for i, key in enumerate(node.obj.keys()):
+            for key in node.obj:
                 match = node.__class__(
                     filter_context=node.filter_context(),
                     obj=key,
                     parent=node,
                     parts=node.parts + (f"{self.env.keys_selector_token}{key}",),
-                    path=f"{node.path}[{self.env.keys_selector_token}][{i}]",
+                    path=f"{node.path}[{self.env.keys_selector_token}{canonical_string(key)}]",
                     root=node.root,
                 )
                 node.add_child(match)
@@ -449,6 +496,98 @@ async def resolve_async(self, node: JSONPathMatch) -> AsyncIterable[JSONPathMatc
                     yield match
 
 
+class KeysFilter(JSONPathSelector):
+    """Selects names from an object’s name/value members.
+
+    NOTE: This is a non-standard selector.
+
+    See https://jg-rp.github.io/json-p3/guides/jsonpath-extra#keys-filter-selector
+    """
+
+    __slots__ = ("expression",)
+
+    def __init__(
+        self,
+        *,
+        env: JSONPathEnvironment,
+        token: Token,
+        expression: BooleanExpression,
+    ) -> None:
+        super().__init__(env=env, token=token)
+        self.expression = expression
+
+    def __str__(self) -> str:
+        return f"~?{self.expression}"
+
+    def __eq__(self, __value: object) -> bool:
+        return (
+            isinstance(__value, Filter)
+            and self.expression == __value.expression
+            and self.token == __value.token
+        )
+
+    def __hash__(self) -> int:
+        return hash(("~", str(self.expression), self.token))
+
+    def resolve(self, node: JSONPathMatch) -> Iterable[JSONPathMatch]:
+        if isinstance(node.value, Mapping):
+            for key, val in node.value.items():
+                context = FilterContext(
+                    env=self.env,
+                    current=val,
+                    root=node.root,
+                    extra_context=node.filter_context(),
+                    current_key=key,
+                )
+
+                try:
+                    if self.expression.evaluate(context):
+                        match = node.__class__(
+                            filter_context=node.filter_context(),
+                            obj=key,
+                            parent=node,
+                            parts=node.parts
+                            + (f"{self.env.keys_selector_token}{key}",),
+                            path=f"{node.path}[{self.env.keys_selector_token}{canonical_string(key)}]",
+                            root=node.root,
+                        )
+                        node.add_child(match)
+                        yield match
+                except JSONPathTypeError as err:
+                    if not err.token:
+                        err.token = self.token
+                    raise
+
+    async def resolve_async(self, node: JSONPathMatch) -> AsyncIterable[JSONPathMatch]:
+        if isinstance(node.value, Mapping):
+            for key, val in node.value.items():
+                context = FilterContext(
+                    env=self.env,
+                    current=val,
+                    root=node.root,
+                    extra_context=node.filter_context(),
+                    current_key=key,
+                )
+
+                try:
+                    if await self.expression.evaluate_async(context):
+                        match = node.__class__(
+                            filter_context=node.filter_context(),
+                            obj=key,
+                            parent=node,
+                            parts=node.parts
+                            + (f"{self.env.keys_selector_token}{key}",),
+                            path=f"{node.path}[{self.env.keys_selector_token}{canonical_string(key)}]",
+                            root=node.root,
+                        )
+                        node.add_child(match)
+                        yield match
+                except JSONPathTypeError as err:
+                    if not err.token:
+                        err.token = self.token
+                    raise
+
+
 class FilterContext:
     """Contextual information and data for evaluating a filter expression."""
 

jsonpath/token.py

@@ -24,6 +24,8 @@
 TOKEN_WILD = sys.intern("TOKEN_WILD")
 TOKEN_NAME = sys.intern("TOKEN_NAME")
 TOKEN_DOT_PROPERTY = sys.intern("TOKEN_DOT_PROPERTY")
+TOKEN_DOT_KEY_PROPERTY = sys.intern("TOKEN_DOT_KEY_PROPERTY")
+TOKEN_KEY_NAME = sys.intern("TOKEN_KEY_NAME")
 
 # Filter expression tokens
 TOKEN_AND = sys.intern("TOKEN_AND")