def test_findall_with_visitors(self) -> None: # Find all assignments in a tree class TestVisitor(m.MatcherDecoratableVisitor): METADATA_DEPENDENCIES: Sequence[meta.ProviderT] = ( meta.ExpressionContextProvider, ) def __init__(self) -> None: super().__init__() self.results: Sequence[cst.CSTNode] = () def visit_Module(self, node: cst.Module) -> None: self.results = self.findall( node, m.MatchMetadata( meta.ExpressionContextProvider, meta.ExpressionContext.STORE ), ) code = """ a = 1 b = True def foo(bar: int) -> bool: return False """ module = cst.parse_module(dedent(code)) wrapper = meta.MetadataWrapper(module) visitor = TestVisitor() wrapper.visit(visitor) self.assertNodeSequenceEqual(visitor.results, [cst.Name("a"), cst.Name("b")])
def test_findall_with_metadata_wrapper(self) -> None: # Find all assignments in a tree code = """ a = 1 b = True def foo(bar: int) -> bool: return False """ module = cst.parse_module(dedent(code)) wrapper = meta.MetadataWrapper(module) # Test that when we find over a wrapper, we implicitly use it for # metadata as well as traversal. booleans = findall( wrapper, m.MatchMetadata(meta.ExpressionContextProvider, meta.ExpressionContext.STORE), ) self.assertNodeSequenceEqual( booleans, [ cst.Name("a"), cst.Name("b"), cst.Name("foo"), cst.Name("bar"), ], ) # Test that we can provide an explicit resolver and tree booleans = findall( wrapper.module, m.MatchMetadata(meta.ExpressionContextProvider, meta.ExpressionContext.STORE), metadata_resolver=wrapper, ) self.assertNodeSequenceEqual( booleans, [ cst.Name("a"), cst.Name("b"), cst.Name("foo"), cst.Name("bar"), ], ) # Test that failing to provide metadata leads to no match booleans = findall( wrapper.module, m.MatchMetadata(meta.ExpressionContextProvider, meta.ExpressionContext.STORE), ) self.assertNodeSequenceEqual(booleans, [])