class UnhashableListTransformer(NoqaAwareTransformer):
@m.call_if_inside(
m.Call(
func=m.OneOf(m.Name(value="set"), m.Name(value="frozenset")),
args=[m.Arg(value=m.OneOf(m.List(), m.Tuple(), m.Set()))],
)
| m.Set() # noqa: W503
)
@m.leave(m.List() | m.Set() | m.Tuple())
def convert_list_arg(
self, _, updated_node: Union[cst.Set, cst.List,
cst.Tuple]) -> cst.BaseExpression:
modified_elements = convert_lists_to_tuples(updated_node.elements)
return updated_node.with_changes(elements=modified_elements)
def test_predicate_logic_on_attributes(self) -> None:
# Verify that we can or things together.
matcher = m.BinaryOperation(left=m.Name(metadata=m.OneOf(
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 1)),
),
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 2)),
),
)))
node, wrapper = self._make_fixture("a + b")
self.assertTrue(matches(node, matcher, metadata_resolver=wrapper))
matcher = m.BinaryOperation(left=m.Integer(metadata=m.OneOf(
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 1)),
),
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 2)),
),
)))
node, wrapper = self._make_fixture("12 + 3")
self.assertTrue(matches(node, matcher, metadata_resolver=wrapper))
node, wrapper = self._make_fixture("123 + 4")
self.assertFalse(matches(node, matcher, metadata_resolver=wrapper))
# Verify that we can and things together
matcher = m.BinaryOperation(left=m.Name(metadata=m.AllOf(
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 1)),
),
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.LOAD),
)))
node, wrapper = self._make_fixture("a + b")
self.assertTrue(matches(node, matcher, metadata_resolver=wrapper))
node, wrapper = self._make_fixture("ab + cd")
self.assertFalse(matches(node, matcher, metadata_resolver=wrapper))
# Verify that we can not things
matcher = m.BinaryOperation(left=m.Name(metadata=m.DoesNotMatch(
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.STORE))))
node, wrapper = self._make_fixture("a + b")
self.assertTrue(matches(node, matcher, metadata_resolver=wrapper))
class HttpRequestXReadLinesTransformer(BaseDjCodemodTransformer):
"""Replace `HttpRequest.xreadlines()` by iterating over the request."""
deprecated_in = DJANGO_2_0
removed_in = DJANGO_3_0
# This should be conservative and only apply changes to:
# - variables called `request`/`req`
# - `request`/`req` attributes (e.g `self.request`/`view.req`...)
matcher = m.Call(
func=m.Attribute(
value=m.OneOf(
m.Name(value="request"),
m.Name(value="req"),
m.Attribute(attr=m.Name(value="request")),
m.Attribute(attr=m.Name(value="req")),
),
attr=m.Name(value="xreadlines"),
)
)
def leave_Call(self, original_node: Call, updated_node: Call) -> BaseExpression:
if not isinstance(updated_node.func, Attribute):
# A bit redundant with matcher below,
# but this is to make type checker happy
return super().leave_Call(original_node, updated_node)
if m.matches(updated_node, self.matcher):
return updated_node.func.value
return super().leave_Call(original_node, updated_node)
def decorator_matcher(self):
matchers_list = self.context.scratch.get(self.ctx_key_decorator_matchers, [])
if len(matchers_list) == 0:
return None
if len(matchers_list) == 1:
return matchers_list[0]
return m.OneOf(*[matcher for matcher in matchers_list])
def visit_Assign(self, node: Assign) -> Optional[bool]:
"""Record variable name the `Library()` call is assigned to."""
if self.library_call_matcher and m.matches(
node,
m.Assign(value=self.library_call_matcher),
):
# Visiting a `register = template.Library()` statement
# Get all names on the left side of the assignment
target_names = (
assign_target.target.value for assign_target in node.targets
)
# Build the decorator matchers to look out for
target_matchers = (
m.Decorator(
decorator=m.Attribute(
value=m.Name(name),
attr=m.Name("assignment_tag"),
)
)
for name in target_names
)
# The final matcher should match if any of the decorator matchers matches
self.context.scratch[self.ctx_key_decorator_matcher] = m.OneOf(
*target_matchers
)
return super().visit_Assign(node)
def __extract_names_multi_assign(self, elements):
# Add self vars. in tuple assignments, e.g. self.x, self.y = 1, 2
# Adds variables in tuple(s) in multiple assignments, e.g. a, (b, c) = 1, (2, 3)
names: List[cst.Name] = []
i = 0
while i < len(elements):
if match.matches(
elements[i],
match.Element(value=match.Name(value=match.DoNotCare()))):
names.append(elements[i].value)
elif match.matches(
elements[i],
match.Element(value=match.Attribute(attr=match.Name(
value=match.DoNotCare())))):
names.append(elements[i].value)
elif match.matches(
elements[i],
match.Element(value=match.Tuple(
elements=match.DoNotCare()))):
elements.extend(
match.findall(
elements[i].value,
match.Element(value=match.OneOf(
match.Attribute(attr=match.Name(
value=match.DoNotCare())),
match.Name(value=match.DoNotCare())))))
i += 1
return names
def __extract_variable_name_type(self, node: cst.AnnAssign):
"""
Extracts a variable's identifier name and its type annotation
"""
return match.extract(
node,
match.AnnAssign( # Annotated Assignment
target=match.OneOf(
match.Name( # Variable name of assignment (only one)
value=match.SaveMatchedNode( # Save result
match.MatchRegex(
r'(.)+'), # Match any string literal
"name")),
# This extracts variables inside __init__ which typically starts with self (e.g. self.x:int=2)
match.Attribute(
value=match.Name(value=match.SaveMatchedNode(
match.MatchRegex(r'(.)+'),
"obj_name" # Object name
)),
attr=match.Name(
match.SaveMatchedNode(match.MatchRegex(r'(.)+'),
"name")),
)),
annotation=match.SaveMatchedNode( # Save result
match.DoNotCare(), # Match any string literal
"type")))
def test_at_least_n_matcher_args_true(self) -> None:
# Match a function call to "foo" where the first argument is the integer
# value 1, and there are at least two wildcard arguments after.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")), m.AtLeastN(m.Arg(), n=2)),
),
)
)
# Match a function call to "foo" where the first argument is the integer
# value 1, and there are at least two arguements are integers of any value
# after.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")), m.AtLeastN(m.Arg(m.Integer()), n=2)),
),
)
)
# Match a function call to "foo" where the first argument is the integer
# value 1, and there are at least two arguements that are integers with the
# value 2 or 3 after.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(
m.Arg(m.Integer("1")),
m.AtLeastN(m.Arg(m.OneOf(m.Integer("2"), m.Integer("3"))), n=2),
),
),
)
)
def test_does_not_match_true(self) -> None:
# Match on any call that takes one argument that isn't the value None.
self.assertTrue(
matches(
libcst.Call(libcst.Name("foo"),
(libcst.Arg(libcst.Name("True")), )),
m.Call(args=(m.Arg(value=m.DoesNotMatch(m.Name("None"))), )),
))
self.assertTrue(
matches(
libcst.Call(libcst.Name("foo"),
(libcst.Arg(libcst.Integer("1")), )),
m.Call(args=(m.DoesNotMatch(m.Arg(m.Name("None"))), )),
))
self.assertTrue(
matches(
libcst.Call(libcst.Name("foo"),
(libcst.Arg(libcst.Integer("1")), )),
m.Call(args=m.DoesNotMatch((m.Arg(m.Integer("2")), ))),
))
# Match any call that takes an argument which isn't True or False.
self.assertTrue(
matches(
libcst.Call(libcst.Name("foo"),
(libcst.Arg(libcst.Integer("1")), )),
m.Call(args=(m.Arg(value=m.DoesNotMatch(
m.OneOf(m.Name("True"), m.Name("False")))), )),
))
# Match any name node that doesn't match the regex for True
self.assertTrue(
matches(
libcst.Name("False"),
m.Name(value=m.DoesNotMatch(m.MatchRegex(r"True"))),
))
def leave_Call(self, original_node: Call,
updated_node: Call) -> BaseExpression:
"""
Remove the `weak` argument if present in the call.
This is only changing calls with keyword arguments.
"""
if self.disconnect_call_matchers and m.matches(
updated_node, m.OneOf(*self.disconnect_call_matchers)):
updated_args = []
should_change = False
last_comma = MaybeSentinel.DEFAULT
# Keep all arguments except the one with the keyword `weak` (if present)
for index, arg in enumerate(updated_node.args):
if m.matches(arg, m.Arg(keyword=m.Name("weak"))):
# An argument with the keyword `weak` was found
# -> we need to rewrite the statement
should_change = True
else:
updated_args.append(arg)
last_comma = arg.comma
if should_change:
# Make sure the end of line is formatted as initially
updated_args[-1] = updated_args[-1].with_changes(
comma=last_comma)
return updated_node.with_changes(args=updated_args)
return super().leave_Call(original_node, updated_node)
def visit_Annotation(self, node: cst.Annotation):
if not match.matches(
node,
match.Annotation(
annotation=match.OneOf(match.Name(
value='None'), match.List(elements=[])))):
self.type_annot_visited = True
def multi(*args, **kwargs):
"""Return a combined matcher for multiple similar types.
*args are the matcher types, and **kwargs are arguments that will be passed to
each type. Returns m.OneOf(...) the results.
"""
return m.OneOf(*(a(**kwargs) for a in args))
def test_at_most_n_matcher_args_true(self) -> None:
# Match a function call to "foo" with at most two arguments, both of which
# are the integer 1.
self.assertTrue(
matches(
cst.Call(func=cst.Name("foo"),
args=(cst.Arg(cst.Integer("1")), )),
m.Call(func=m.Name("foo"),
args=(m.AtMostN(m.Arg(m.Integer("1")), n=2), )),
))
# Match a function call to "foo" with at most two arguments, both of which
# can be the integer 1 or 2.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2"))),
),
m.Call(
func=m.Name("foo"),
args=(m.AtMostN(m.Arg(
m.OneOf(m.Integer("1"), m.Integer("2"))),
n=2), ),
),
))
# Match a function call to "foo" with at most two arguments, the first
# one being the integer 1 and the second one, if included, being the
# integer 2.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2"))),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")),
m.ZeroOrOne(m.Arg(m.Integer("2")))),
),
))
# Match a function call to "foo" with at most six arguments, the first
# one being the integer 1 and the second one, if included, being the
# integer 2.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2"))),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")),
m.ZeroOrOne(m.Arg(m.Integer("2")))),
),
))
def test_or_matcher_true(self) -> None:
# Match on either True or False identifier.
self.assertTrue(
matches(libcst.Name("True"), m.OneOf(m.Name("True"), m.Name("False")))
)
# Match any assignment that assigns a value of True or False to an
# unspecified target.
self.assertTrue(
matches(
libcst.Assign(
(libcst.AssignTarget(libcst.Name("x")),), libcst.Name("True")
),
m.Assign(value=m.OneOf(m.Name("True"), m.Name("False"))),
)
)
self.assertTrue(
matches(
libcst.Call(
libcst.Name("foo"),
(
libcst.Arg(libcst.Integer("1")),
libcst.Arg(libcst.Integer("2")),
libcst.Arg(libcst.Integer("3")),
),
),
m.Call(
m.Name("foo"),
m.OneOf(
(
m.Arg(m.Integer("3")),
m.Arg(m.Integer("2")),
m.Arg(m.Integer("1")),
),
(
m.Arg(m.Integer("1")),
m.Arg(m.Integer("2")),
m.Arg(m.Integer("3")),
),
),
),
)
)
def test_or_matcher_false(self) -> None:
# Fail to match since None is not True or False.
self.assertFalse(
matches(libcst.Name("None"), m.OneOf(m.Name("True"), m.Name("False")))
)
# Fail to match since assigning None to a target is not the same as
# assigning True or False to a target.
self.assertFalse(
matches(
libcst.Assign(
(libcst.AssignTarget(libcst.Name("x")),), libcst.Name("None")
),
m.Assign(value=m.OneOf(m.Name("True"), m.Name("False"))),
)
)
self.assertFalse(
matches(
libcst.Call(
libcst.Name("foo"),
(
libcst.Arg(libcst.Integer("1")),
libcst.Arg(libcst.Integer("2")),
libcst.Arg(libcst.Integer("3")),
),
),
m.Call(
m.Name("foo"),
m.OneOf(
(
m.Arg(m.Integer("3")),
m.Arg(m.Integer("2")),
m.Arg(m.Integer("1")),
),
(
m.Arg(m.Integer("4")),
m.Arg(m.Integer("5")),
m.Arg(m.Integer("6")),
),
),
),
)
)
def leave_Expr(
self, original_node: "Expr", updated_node: "Expr"
) -> Union["BaseSmallStatement", RemovalSentinel]:
# FIXME: For some strange reason if we put that matcher combination
# into m.call_if_inside() (or others), then it will get triggered on
# _every_ function call. Not sure if bug or a feature :/
if m.matches(
original_node,
m.Expr(value=m.Call(func=m.OneOf(
m.Attribute(
value=m.OneOf(m.Name(value="pdb"), m.Name(
value="ipdb")),
attr=m.Name(value="set_trace"),
),
m.Name("breakpoint"),
))),
):
return cst.RemoveFromParent()
return original_node
def match_calls_str_format(node: cst.Call) -> bool:
# Check if a call to str.format() is being made
return m.matches(
node,
m.Call(
func=m.Attribute(
value=m.OneOf(m.SimpleString(), m.ConcatenatedString()),
attr=m.Name("format"),
)
),
)
def visit_BinaryOperation(self, node: cst.BinaryOperation) -> None:
if not self.logging_stack:
return
if m.matches(
node,
m.BinaryOperation(
left=m.OneOf(m.SimpleString(), m.ConcatenatedString()),
operator=m.Modulo(),
),
):
self.report(node)
def get_import_name_from_attr(attr_node: cst.Attribute) -> str:
name = [attr_node.attr.value] # last value
node = attr_node.value
while m.matches(node, m.OneOf(m.Name(), m.Attribute())):
if isinstance(node, cst.Attribute):
name.append(node.attr.value)
node = node.value
else:
name.append(cst.ensure_type(node, cst.Name).value)
break
name.reverse()
return ".".join(name)
def __get_var_names_counter(self, node, scope):
vars_name = match.extractall(
node,
match.OneOf(
match.AssignTarget(target=match.SaveMatchedNode(
match.Name(value=match.DoNotCare()), "name")),
match.AnnAssign(target=match.SaveMatchedNode(
match.Name(value=match.DoNotCare()), "name"))))
return Counter([
n['name'].value for n in vars_name if isinstance(
self.get_metadata(cst.metadata.ScopeProvider, n['name']),
scope)
])
def visit_Assign(self, node: Assign) -> Optional[bool]:
"""Record variable name the `Library()` call is assigned to."""
if self.library_call_matcher and m.matches(
node,
m.Assign(value=self.library_call_matcher),
):
# Visiting a `register = template.Library()` statement
# Generate decorator matchers based on left hand side names
decorator_matchers = self._gen_decorator_matchers(node.targets)
# should match if any of the decorator matches
self.context.scratch[self.ctx_key_decorator_matcher] = m.OneOf(
*decorator_matchers
)
return super().visit_Assign(node)
def visit_Call(self, node: cst.Call) -> None:
# print(node)
d = m.extract(
node,
m.Call(
func=m.OneOf(m.Name("Extension"), m.Name("addMacExtension")),
args=(
m.Arg(value=m.SaveMatchedNode(m.SimpleString(),
"extension_name")),
m.ZeroOrMore(m.DoNotCare()),
),
),
)
if d:
assert isinstance(d["extension_name"], cst.SimpleString)
self.extension_names.append(d["extension_name"].evaluated_value)
def contains_union_with_none(self, node: cst.Annotation) -> bool:
return m.matches(
node,
m.Annotation(
m.Subscript(
value=m.Name("Union"),
slice=m.OneOf(
[
m.SubscriptElement(m.Index()),
m.SubscriptElement(m.Index(m.Name("None"))),
],
[
m.SubscriptElement(m.Index(m.Name("None"))),
m.SubscriptElement(m.Index()),
],
),
)),
)
def leave_ImportFrom(
self, original_node: ImportFrom, updated_node: ImportFrom
) -> Union[BaseSmallStatement, RemovalSentinel]:
base_cls_matcher = []
if m.matches(
updated_node,
m.ImportFrom(module=module_matcher(["django", "contrib", "admin"])),
):
for imported_name in updated_node.names:
if m.matches(
imported_name, m.ImportAlias(name=m.Name("TabularInline"))
):
base_cls_matcher.append(m.Arg(m.Name("TabularInline")))
if m.matches(
imported_name, m.ImportAlias(name=m.Name("StackedInline"))
):
base_cls_matcher.append(m.Arg(m.Name("StackedInline")))
if m.matches(
updated_node,
m.ImportFrom(module=module_matcher(["django", "contrib"])),
):
for imported_name in updated_node.names:
if m.matches(imported_name, m.ImportAlias(name=m.Name("admin"))):
base_cls_matcher.extend(
[
m.Arg(
m.Attribute(
value=m.Name("admin"), attr=m.Name("TabularInline")
)
),
m.Arg(
m.Attribute(
value=m.Name("admin"), attr=m.Name("StackedInline")
)
),
]
)
# Save valid matchers in the context
if base_cls_matcher:
self.context.scratch[self.ctx_key_base_cls_matcher] = m.OneOf(
*base_cls_matcher
)
return super().leave_ImportFrom(original_node, updated_node)
class InlineHasAddPermissionsTransformer(ContextAwareTransformer):
"""Add the ``obj`` argument to ``InlineModelAdmin.has_add_permission()``."""
context_key = "InlineHasAddPermissionsTransformer"
base_cls_matcher = m.OneOf(
m.Arg(m.Attribute(value=m.Name("admin"),
attr=m.Name("TabularInline"))),
m.Arg(m.Name("TabularInline")),
m.Arg(m.Attribute(value=m.Name("admin"),
attr=m.Name("StackedInline"))),
m.Arg(m.Name("StackedInline")),
)
def visit_ClassDef_bases(self, node: ClassDef) -> None:
for base_cls in node.bases:
if m.matches(base_cls, self.base_cls_matcher):
self.context.scratch[self.context_key] = True
super().visit_ClassDef_bases(node)
def leave_ClassDef(
self, original_node: ClassDef,
updated_node: ClassDef) -> Union[BaseStatement, RemovalSentinel]:
self.context.scratch.pop(self.context_key, None)
return super().leave_ClassDef(original_node, updated_node)
@property
def _is_context_right(self):
return self.context.scratch.get(self.context_key, False)
def leave_FunctionDef(
self, original_node: FunctionDef, updated_node: FunctionDef
) -> Union[BaseStatement, RemovalSentinel]:
if (m.matches(updated_node,
m.FunctionDef(name=m.Name("has_add_permission")))
and self._is_context_right):
if len(updated_node.params.params) == 2:
old_params = updated_node.params
updated_params = old_params.with_changes(params=(
*old_params.params,
Param(name=Name("obj"), default=Name("None")),
))
return updated_node.with_changes(params=updated_params)
return super().leave_FunctionDef(original_node, updated_node)
def _has_testnode(node: cst.Module) -> bool:
return m.matches(
node,
m.Module(body=[
# Sequence wildcard matchers matches LibCAST nodes in a row in a
# sequence. It does not implicitly match on partial sequences. So,
# when matching against a sequence we will need to provide a
# complete pattern. This often means using helpers such as
# ``ZeroOrMore()`` as the first and last element of the sequence.
m.ZeroOrMore(),
m.AtLeastN(
n=1,
matcher=m.OneOf(
m.FunctionDef(name=m.Name(value=m.MatchIfTrue(
lambda value: value.startswith("test_")))),
m.ClassDef(name=m.Name(value=m.MatchIfTrue(
lambda value: value.startswith("Test")))),
),
),
m.ZeroOrMore(),
]),
)
def __extract_variable_name(self, node: cst.AssignTarget):
extracted_var_names = match.extract(
node,
match.AssignTarget( # Assignment operator
target=match.OneOf( # Two cases exist
match.Name( # Single target
value=match.SaveMatchedNode( # Save result
match.MatchRegex(
r'(.)+'), # Match any string literal
"name")),
match.Tuple( # Multi-target
elements=match.SaveMatchedNode( # Save result
match.DoNotCare(), # Type of list
"names")),
# This extracts variables inside __init__ without type annotation (e.g. self.x=2)
match.Attribute(
value=match.Name(value=match.SaveMatchedNode(
match.MatchRegex(r'(.)+'),
"obj_name" # Object name
)),
attr=match.Name(
match.SaveMatchedNode(match.MatchRegex(r'(.)+'),
"name")),
))))
if extracted_var_names is not None:
if "name" in extracted_var_names:
t = self.__get_type_from_metadata(node.target)
extracted_var_names['type'] = (t, INF_TYPE_ANNOT
if t else UNK_TYPE_ANNOT)
return extracted_var_names
elif "names" in extracted_var_names:
return {
'names':
self.__extract_names_multi_assign(
list(extracted_var_names['names']))
}
else:
return extracted_var_names
def some_version_of(tag: str) -> m.OneOf[m.Union[m.Name, m.Attribute]]:
"""
Poorly named wrapper around name_attr_possibilities.
"""
return m.OneOf(*name_attr_possibilities(tag))
def test_zero_or_more_matcher_args_true(self) -> None:
# Match a function call to "foo" where the first argument is the integer
# value 1, and the rest of the arguements are wildcards.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")), m.ZeroOrMore(m.Arg())),
),
)
)
# Match a function call to "foo" where the first argument is the integer
# value 1, and the rest of the arguements are integers of any value.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(m.Arg(m.Integer("1")), m.ZeroOrMore(m.Arg(m.Integer()))),
),
)
)
# Match a function call to "foo" with zero or more arguments, where the
# first argument can optionally be the integer 1 or 2, and the second
# can only be the integer 2. This case verifies non-greedy behavior in the
# matcher.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(
m.ZeroOrMore(m.Arg(m.OneOf(m.Integer("1"), m.Integer("2")))),
m.Arg(m.Integer("2")),
m.ZeroOrMore(),
),
),
)
)
# Match a function call to "foo" where the first argument is the integer
# value 1, and the rest of the arguements are integers with the value
# 2 or 3.
self.assertTrue(
matches(
cst.Call(
func=cst.Name("foo"),
args=(
cst.Arg(cst.Integer("1")),
cst.Arg(cst.Integer("2")),
cst.Arg(cst.Integer("3")),
),
),
m.Call(
func=m.Name("foo"),
args=(
m.Arg(m.Integer("1")),
m.ZeroOrMore(m.Arg(m.OneOf(m.Integer("2"), m.Integer("3")))),
),
),
)
)
def oneof_names(*names):
return m.OneOf(*map(m.Name, names))
