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, [])