diff --git a/docs/source/matchers.rst b/docs/source/matchers.rst index 210ae3ac..6bdf214e 100644 --- a/docs/source/matchers.rst +++ b/docs/source/matchers.rst @@ -32,6 +32,7 @@ selectively control when LibCST calls visitor functions. .. autofunction:: libcst.matchers.matches .. autofunction:: libcst.matchers.findall .. autofunction:: libcst.matchers.extract +.. autofunction:: libcst.matchers.extractall .. _libcst-matcher-decorators: diff --git a/libcst/codegen/gen_matcher_classes.py b/libcst/codegen/gen_matcher_classes.py index f2bfac47..adb3aaad 100644 --- a/libcst/codegen/gen_matcher_classes.py +++ b/libcst/codegen/gen_matcher_classes.py @@ -482,7 +482,7 @@ generated_code.append("from typing_extensions import Literal") generated_code.append("import libcst as cst") generated_code.append("") generated_code.append( - "from libcst.matchers._matcher_base import BaseMatcherNode, DoNotCareSentinel, DoNotCare, OneOf, AllOf, DoesNotMatch, MatchIfTrue, MatchRegex, MatchMetadata, MatchMetadataIfTrue, ZeroOrMore, AtLeastN, ZeroOrOne, AtMostN, SaveMatchedNode, extract, findall, matches" + "from libcst.matchers._matcher_base import BaseMatcherNode, DoNotCareSentinel, DoNotCare, OneOf, AllOf, DoesNotMatch, MatchIfTrue, MatchRegex, MatchMetadata, MatchMetadataIfTrue, ZeroOrMore, AtLeastN, ZeroOrOne, AtMostN, SaveMatchedNode, extract, extractall, findall, matches" ) all_exports.update( [ @@ -502,6 +502,7 @@ all_exports.update( "AtMostN", "SaveMatchedNode", "extract", + "extractall", "findall", "matches", ] diff --git a/libcst/matchers/__init__.py b/libcst/matchers/__init__.py index cf12960a..aa4ef486 100644 --- a/libcst/matchers/__init__.py +++ b/libcst/matchers/__init__.py @@ -31,6 +31,7 @@ from libcst.matchers._matcher_base import ( ZeroOrMore, ZeroOrOne, extract, + extractall, findall, matches, ) @@ -13466,6 +13467,7 @@ __all__ = [ "call_if_inside", "call_if_not_inside", "extract", + "extractall", "findall", "leave", "matches", diff --git a/libcst/matchers/_matcher_base.py b/libcst/matchers/_matcher_base.py index 7b6197ea..dedf9ca7 100644 --- a/libcst/matchers/_matcher_base.py +++ b/libcst/matchers/_matcher_base.py @@ -18,6 +18,7 @@ from typing import ( Optional, Pattern, Sequence, + Tuple, Type, TypeVar, Union, @@ -296,15 +297,15 @@ class _InverseOf(Generic[_MatcherT]): class _ExtractMatchingNode(Generic[_MatcherT]): """ Transparent pass-through matcher that captures the node which matches its children, - making it available to the caller of :func:`extract`. + making it available to the caller of :func:`extract` or :func:`extractall`. Note that you should refrain from constructing a :class:`_ExtractMatchingNode` directly, and should instead use the :func:`SaveMatchedNode` helper function. For example, the following will match against any binary operation whose left and right operands are not integers, saving those expressions for later inspection. - If used inside a :func:`extract`, the resulting dictionary will contain the - keys ``left_operand`` and ``right_operand``. + If used inside :func:`extract` or :func:`extractall`, the resulting dictionary will + contain the keys ``left_operand`` and ``right_operand``. m.BinaryOperation( left=m.SaveMatchedNode( @@ -335,7 +336,7 @@ class _ExtractMatchingNode(Generic[_MatcherT]): def name(self) -> str: """ The name we will call our captured LibCST node inside the resulting dictionary - returned by :func:`extract`. + returned by :func:`extract` or :func:`extractall`. """ return self._name @@ -904,12 +905,13 @@ def DoesNotMatch(obj: _OtherNodeT) -> _OtherNodeT: def SaveMatchedNode(matcher: _OtherNodeT, name: str) -> _OtherNodeT: """ Matcher helper that captures the matched node that matched against a matcher - class, making it available in the dictionary returned by :func:`extract`. + class, making it available in the dictionary returned by :func:`extract` or + :func:`extractall`. For example, the following will match against any binary operation whose left and right operands are not integers, saving those expressions for later inspection. - If used inside a :func:`extract`, the resulting dictionary will contain the - keys ``left_operand`` and ``right_operand``:: + If used inside :func:`extract` or :func:`extractall`, the resulting dictionary + will contain the keys ``left_operand`` and ``right_operand``:: m.BinaryOperation( left=m.SaveMatchedNode( @@ -1472,19 +1474,30 @@ class _FindAllVisitor(libcst.CSTVisitor): self.matcher = matcher self.metadata_lookup = metadata_lookup self.found_nodes: List[libcst.CSTNode] = [] + self.extracted_nodes: List[ + Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]] + ] = [] def on_visit(self, node: libcst.CSTNode) -> bool: - if _matches(node, self.matcher, self.metadata_lookup) is not None: + match = _matches(node, self.matcher, self.metadata_lookup) + if match is not None: self.found_nodes.append(node) + self.extracted_nodes.append(match) return True -def findall( +def _find_or_extract_all( tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper], matcher: Union[ BaseMatcherNode, MatchIfTrue[Callable[[object], bool]], _BaseMetadataMatcher, + # The inverse clause is left off of the public functions `findall` and + # `extractall` because we play a dirty trick. We lie to the typechecker + # that `DoesNotMatch` returns identity, so the public functions don't + # need to be aware of inverses. If we could represent predicate logic + # in python types we could get away with this, but that's not the state + # of things right now. _InverseOf[ Union[ BaseMatcherNode, @@ -1497,11 +1510,49 @@ def findall( metadata_resolver: Optional[ Union[libcst.MetadataDependent, libcst.MetadataWrapper] ] = None, +) -> Tuple[ + Sequence[libcst.CSTNode], + Sequence[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]], +]: + if isinstance(tree, (RemovalSentinel, MaybeSentinel)): + # We can't possibly match on a removal sentinel, so it doesn't match. + return [], [] + if isinstance(matcher, (AtLeastN, AtMostN)): + # We can't match this, since these matchers are forbidden at top level. + # These are not subclasses of BaseMatcherNode, but in the case that the + # user is not using type checking, this should still behave correctly. + return [], [] + + if isinstance(tree, meta.MetadataWrapper) and metadata_resolver is None: + # Provide a convenience for calling findall directly on a MetadataWrapper. + metadata_resolver = tree + + if metadata_resolver is None: + fetcher = _construct_metadata_fetcher_null() + elif isinstance(metadata_resolver, libcst.MetadataWrapper): + fetcher = _construct_metadata_fetcher_wrapper(metadata_resolver) + else: + fetcher = _construct_metadata_fetcher_dependent(metadata_resolver) + + finder = _FindAllVisitor(matcher, fetcher) + tree.visit(finder) + return finder.found_nodes, finder.extracted_nodes + + +def findall( + tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper], + matcher: Union[ + BaseMatcherNode, MatchIfTrue[Callable[[object], bool]], _BaseMetadataMatcher, + ], + *, + metadata_resolver: Optional[ + Union[libcst.MetadataDependent, libcst.MetadataWrapper] + ] = None, ) -> Sequence[libcst.CSTNode]: """ - Given an arbitrary node from a LibCST tree, and an arbitrary matcher, iterates - over that node and all children, returning a sequence of all child nodes that - match the given matcher. Note that the node can also be a + Given an arbitrary node from a LibCST tree and an arbitrary matcher, iterates + over that node and all children returning a sequence of all child nodes that + match the given matcher. Note that the tree can also be a :class:`~libcst.RemovalSentinel` or a :class:`~libcst.MaybeSentinel` in order to use findall directly on transform results and node attributes. In these cases, :func:`findall` will always return an empty sequence. Note also that @@ -1519,26 +1570,44 @@ def findall( :class:`AtMostN` matcher because these types are wildcards which can only be usedi inside sequences. """ - if isinstance(tree, (RemovalSentinel, MaybeSentinel)): - # We can't possibly match on a removal sentinel, so it doesn't match. - return [] - if isinstance(matcher, (AtLeastN, AtMostN)): - # We can't match this, since these matchers are forbidden at top level. - # These are not subclasses of BaseMatcherNode, but in the case that the - # user is not using type checking, this should still behave correctly. - return [] + nodes, _ = _find_or_extract_all(tree, matcher, metadata_resolver=metadata_resolver) + return nodes - if isinstance(tree, meta.MetadataWrapper) and metadata_resolver is None: - # Provide a convenience for calling findall directly on a MetadataWrapper. - metadata_resolver = tree - if metadata_resolver is None: - fetcher = _construct_metadata_fetcher_null() - elif isinstance(metadata_resolver, libcst.MetadataWrapper): - fetcher = _construct_metadata_fetcher_wrapper(metadata_resolver) - else: - fetcher = _construct_metadata_fetcher_dependent(metadata_resolver) +def extractall( + tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper], + matcher: Union[ + BaseMatcherNode, MatchIfTrue[Callable[[object], bool]], _BaseMetadataMatcher, + ], + *, + metadata_resolver: Optional[ + Union[libcst.MetadataDependent, libcst.MetadataWrapper] + ] = None, +) -> Sequence[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]: + """ + Given an arbitrary node from a LibCST tree and an arbitrary matcher, iterates + over that node and all children returning a sequence of dictionaries representing + the saved and extracted children specified by :func:`SaveMatchedNode` for each + match found in the tree. This is analogous to running a :func:`findall` over a + tree, then running :func:`extract` with the same matcher over each of the returned + nodes. Note that the tree can also be a :class:`~libcst.RemovalSentinel` or a + :class:`~libcst.MaybeSentinel` in order to use extractall directly on transform + results and node attributes. In these cases, :func:`extractall` will always + return an empty sequence. Note also that instead of a LibCST tree, you can + instead pass in a :class:`~libcst.metadata.MetadataWrapper`. This mirrors the + fact that you can call ``visit`` on a :class:`~libcst.metadata.MetadataWrapper` + in order to iterate over it with a transform. If you provide a wrapper for the + tree and do not set the ``metadata_resolver`` parameter specifically, it will + automatically be set to the wrapper for you. - finder = _FindAllVisitor(matcher, fetcher) - tree.visit(finder) - return finder.found_nodes + The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`, + or a :class:`OneOf`/:class:`AllOf` special matcher. Unlike :func:`matches`, it can + also be a :class:`MatchIfTrue` or :func:`DoesNotMatch` matcher, since we are + traversing the tree looking for matches. It cannot be a :class:`AtLeastN` or + :class:`AtMostN` matcher because these types are wildcards which can only be usedi + inside sequences. + """ + _, extractions = _find_or_extract_all( + tree, matcher, metadata_resolver=metadata_resolver + ) + return extractions diff --git a/libcst/matchers/_visitors.py b/libcst/matchers/_visitors.py index adf0f615..92a0d956 100644 --- a/libcst/matchers/_visitors.py +++ b/libcst/matchers/_visitors.py @@ -36,9 +36,10 @@ from libcst.matchers._matcher_base import ( BaseMatcherNode, MatchIfTrue, MatchMetadata, + MatchMetadataIfTrue, OneOf, - _InverseOf, extract, + extractall, findall, matches, ) @@ -565,9 +566,7 @@ class MatcherDecoratableTransformer(CSTTransformer): BaseMatcherNode, MatchIfTrue[Callable[..., bool]], MatchMetadata, - _InverseOf[ - Union[BaseMatcherNode, MatchIfTrue[Callable[..., bool]], MatchMetadata] - ], + MatchMetadataIfTrue, ], ) -> Sequence[cst.CSTNode]: """ @@ -593,6 +592,25 @@ class MatcherDecoratableTransformer(CSTTransformer): """ return extract(node, matcher, metadata_resolver=self) + def extractall( + self, + tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode], + matcher: Union[ + BaseMatcherNode, + MatchIfTrue[Callable[..., bool]], + MatchMetadata, + MatchMetadataIfTrue, + ], + ) -> Sequence[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]: + """ + A convenience method to call :func:`~libcst.matchers.extractall` without requiring + an explicit parameter for metadata. Since our instance is an instance of + :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see + documentation for :func:`~libcst.matchers.extractall` as it is identical to this + function. + """ + return extractall(tree, matcher, metadata_resolver=self) + def _transform_module_impl(self, tree: cst.Module) -> cst.Module: return tree.visit(self) @@ -718,9 +736,7 @@ class MatcherDecoratableVisitor(CSTVisitor): BaseMatcherNode, MatchIfTrue[Callable[..., bool]], MatchMetadata, - _InverseOf[ - Union[BaseMatcherNode, MatchIfTrue[Callable[..., bool]], MatchMetadata] - ], + MatchMetadataIfTrue, ], ) -> Sequence[cst.CSTNode]: """ @@ -745,3 +761,22 @@ class MatcherDecoratableVisitor(CSTVisitor): function. """ return extract(node, matcher, metadata_resolver=self) + + def extractall( + self, + tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode], + matcher: Union[ + BaseMatcherNode, + MatchIfTrue[Callable[..., bool]], + MatchMetadata, + MatchMetadataIfTrue, + ], + ) -> Sequence[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]: + """ + A convenience method to call :func:`~libcst.matchers.extractall` without requiring + an explicit parameter for metadata. Since our instance is an instance of + :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see + documentation for :func:`~libcst.matchers.extractall` as it is identical to this + function. + """ + return extractall(tree, matcher, metadata_resolver=self) diff --git a/libcst/matchers/tests/test_findall.py b/libcst/matchers/tests/test_findall.py index 8b36616a..e57efb0a 100644 --- a/libcst/matchers/tests/test_findall.py +++ b/libcst/matchers/tests/test_findall.py @@ -11,7 +11,7 @@ from typing import Optional, Sequence import libcst as cst import libcst.matchers as m import libcst.metadata as meta -from libcst.matchers import findall +from libcst.matchers import extractall, findall from libcst.testing.utils import UnitTest @@ -162,3 +162,16 @@ class MatchersFindAllTest(UnitTest): visitor = TestTransformer() wrapper.visit(visitor) self.assertNodeSequenceEqual(visitor.results, [cst.Name("a"), cst.Name("b")]) + + +class MatchersExtractAllTest(UnitTest): + def test_extractall_simple(self) -> None: + expression = cst.parse_expression("a + b[c], d(e, f * g, h.i.j)") + matches = extractall(expression, m.Arg(m.SaveMatchedNode(~m.Name(), "expr"))) + extracted_args = cst.ensure_type( + cst.ensure_type(expression, cst.Tuple).elements[1].value, cst.Call, + ).args + self.assertEqual( + matches, + [{"expr": extracted_args[1].value}, {"expr": extracted_args[2].value}], + )