def test_extract_sequence_element(self) -> None:
# Verify true behavior
expression = cst.parse_expression("a + b[c], d(e, f * g, h.i.j)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=[m.SaveMatchedNode(m.ZeroOrMore(), "args")])),
]),
)
extracted_seq = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).args
self.assertEqual(nodes, {"args": extracted_seq})
# Verify false behavior
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=[
m.SaveMatchedNode(m.ZeroOrMore(m.Arg(m.Subscript())),
"args")
])),
]),
)
self.assertIsNone(nodes)
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 test_extract_simple(self) -> None:
# Verify true behavior
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(
left=m.SaveMatchedNode(m.Name(), "left"))),
m.Element(m.Call()),
]),
)
extracted_node = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[0].value,
cst.BinaryOperation,
).left
self.assertEqual(nodes, {"left": extracted_node})
# Verify false behavior
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(
left=m.SaveMatchedNode(m.Subscript(), "left"))),
m.Element(m.Call()),
]),
)
self.assertIsNone(nodes)
def test_extract_tautology(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.SaveMatchedNode(
m.Tuple(elements=[
m.Element(m.BinaryOperation()),
m.Element(m.Call())
]),
name="node",
),
)
self.assertEqual(nodes, {"node": expression})
def test_extract_metadata(self) -> None:
# Verify true behavior
module = cst.parse_module("a + b[c], d(e, f * g)")
wrapper = cst.MetadataWrapper(module)
expression = cst.ensure_type(
cst.ensure_type(wrapper.module.body[0],
cst.SimpleStatementLine).body[0],
cst.Expr,
).value
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(left=m.Name(metadata=m.SaveMatchedNode(
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 1)),
),
"left",
)))),
m.Element(m.Call()),
]),
metadata_resolver=wrapper,
)
extracted_node = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[0].value,
cst.BinaryOperation,
).left
self.assertEqual(nodes, {"left": extracted_node})
# Verify false behavior
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(left=m.Name(metadata=m.SaveMatchedNode(
m.MatchMetadata(
meta.PositionProvider,
self._make_coderange((1, 0), (1, 2)),
),
"left",
)))),
m.Element(m.Call()),
]),
metadata_resolver=wrapper,
)
self.assertIsNone(nodes)
def test_extract_predicates(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(
left=m.SaveMatchedNode(m.Name(), "left"))),
m.Element(
m.Call(func=m.SaveMatchedNode(m.Name(), "func")
| m.SaveMatchedNode(m.Attribute(), "attr"))),
]),
)
extracted_node_left = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[0].value,
cst.BinaryOperation,
).left
extracted_node_func = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).func
self.assertEqual(nodes, {
"left": extracted_node_left,
"func": extracted_node_func
})
expression = cst.parse_expression("a + b[c], d.z(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.Element(
m.BinaryOperation(
left=m.SaveMatchedNode(m.Name(), "left"))),
m.Element(
m.Call(func=m.SaveMatchedNode(m.Name(), "func")
| m.SaveMatchedNode(m.Attribute(), "attr"))),
]),
)
extracted_node_left = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[0].value,
cst.BinaryOperation,
).left
extracted_node_attr = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).func
self.assertEqual(nodes, {
"left": extracted_node_left,
"attr": extracted_node_attr
})
def visit_Element(self, node: cst.Element) -> None:
# We only care about elements in *List* or *Tuple* specifically coming from
# inside the multiple assignments.
if self._assigntarget_counter > 0:
if m.matches(node, m.Element(value=m.Name())):
nodename = cst.ensure_type(node.value, cst.Name).value
self._validate_nodename(node, nodename,
NamingConvention.SNAKE_CASE)
def visit_Element(self, node: libcst.Element) -> bool:
# See if this is a entry that is a string.
extraction = self.extract(
node, m.Element(m.SaveMatchedNode(m.SimpleString(), "string")))
if extraction:
string = ensure_type(extraction["string"], libcst.SimpleString)
self.explicit_exported_objects.add(string.evaluated_value)
# Don't need to visit children
return False
def test_extract_optional_wildcard_tail(self) -> None:
expression = cst.parse_expression("[3]")
nodes = m.extract(
expression,
m.List(elements=[
m.Element(value=m.Integer(value="3")),
m.SaveMatchedNode(m.ZeroOrMore(), "tail1"),
m.SaveMatchedNode(m.ZeroOrMore(), "tail2"),
]),
)
self.assertEqual(nodes, {"tail1": (), "tail2": ()})
def test_extract_sequence(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g, h.i.j)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=m.SaveMatchedNode([m.ZeroOrMore()], "args"))),
]),
)
extracted_seq = cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).args
self.assertEqual(nodes, {"args": extracted_seq})
def test_extract_optional_wildcard(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=[
m.ZeroOrMore(),
m.ZeroOrOne(
m.Arg(m.SaveMatchedNode(m.Attribute(), "arg"))),
])),
]),
)
self.assertEqual(nodes, {})
def convert_lists_to_tuples(
elements: Sequence[cst.BaseElement], ) -> List[cst.BaseElement]:
result: List[cst.BaseElement] = []
for element in elements:
if m.matches(element, m.Element(value=m.List())):
unhashable_list: cst.List = cst.ensure_type(
element.value, cst.List)
result.append(
element.with_changes(value=cst.Tuple(
elements=convert_lists_to_tuples(
unhashable_list.elements))))
else:
result.append(element)
return result
def test_extract_precedence_sequence_wildcard(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=[
m.ZeroOrMore(
m.Arg(m.SaveMatchedNode(m.DoNotCare(), "arg")))
])),
]),
)
extracted_node = (cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).args[1].value)
self.assertEqual(nodes, {"arg": extracted_node})
def test_extract_precedence_parent(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.SaveMatchedNode(
m.Call(args=[
m.Arg(m.SaveMatchedNode(m.Name(), "name")),
m.DoNotCare(),
]),
"name",
)),
]),
)
extracted_node = cst.ensure_type(expression,
cst.Tuple).elements[1].value
self.assertEqual(nodes, {"name": extracted_node})
def test_extract_optional_wildcard_present(self) -> None:
expression = cst.parse_expression("a + b[c], d(e, f * g, h.i.j)")
nodes = m.extract(
expression,
m.Tuple(elements=[
m.DoNotCare(),
m.Element(
m.Call(args=[
m.DoNotCare(),
m.DoNotCare(),
m.ZeroOrOne(
m.Arg(m.SaveMatchedNode(m.Attribute(), "arg"))),
])),
]),
)
extracted_node = (cst.ensure_type(
cst.ensure_type(expression, cst.Tuple).elements[1].value,
cst.Call).args[2].value)
self.assertEqual(nodes, {"arg": extracted_node})
def test_extract_sequence_multiple_wildcards(self) -> None:
expression = cst.parse_expression("1, 2, 3, 4")
nodes = m.extract(
expression,
m.Tuple(elements=(
m.SaveMatchedNode(m.ZeroOrMore(), "head"),
m.SaveMatchedNode(m.Element(value=m.Integer(
value="3")), "element"),
m.SaveMatchedNode(m.ZeroOrMore(), "tail"),
)),
)
tuple_elements = cst.ensure_type(expression, cst.Tuple).elements
self.assertEqual(
nodes,
{
"head": tuple(tuple_elements[:2]),
"element": tuple_elements[2],
"tail": tuple(tuple_elements[3:]),
},
)
def visit_Call(self, node: cst.Call) -> None:
if m.matches(
node,
m.Call(
func=m.Name("tuple") | m.Name("list") | m.Name("set")
| m.Name("dict"),
args=[m.Arg(value=m.List() | m.Tuple())],
),
) or m.matches(
node,
m.Call(func=m.Name("tuple") | m.Name("list") | m.Name("dict"),
args=[]),
):
pairs_matcher = m.ZeroOrMore(
m.Element(m.Tuple(
elements=[m.DoNotCare(), m.DoNotCare()]))
| m.Element(m.List(
elements=[m.DoNotCare(), m.DoNotCare()])))
exp = cst.ensure_type(node, cst.Call)
call_name = cst.ensure_type(exp.func, cst.Name).value
# If this is a empty call, it's an Unnecessary Call where we rewrite the call
# to literal, except set().
if not exp.args:
elements = []
message_formatter = UNNCESSARY_CALL
else:
arg = exp.args[0].value
elements = cst.ensure_type(
arg, cst.List
if isinstance(arg, cst.List) else cst.Tuple).elements
message_formatter = UNNECESSARY_LITERAL
if call_name == "tuple":
new_node = cst.Tuple(elements=elements)
elif call_name == "list":
new_node = cst.List(elements=elements)
elif call_name == "set":
# set() doesn't have an equivelant literal call. If it was
# matched here, it's an unnecessary literal suggestion.
if len(elements) == 0:
self.report(
node,
UNNECESSARY_LITERAL.format(func=call_name),
replacement=node.deep_replace(
node, cst.Call(func=cst.Name("set"))),
)
return
new_node = cst.Set(elements=elements)
elif len(elements) == 0 or m.matches(
exp.args[0].value,
m.Tuple(elements=[pairs_matcher])
| m.List(elements=[pairs_matcher]),
):
new_node = cst.Dict(elements=[(
lambda val: cst.DictElement(val.elements[
0].value, val.elements[1].value))(cst.ensure_type(
ele.value,
cst.Tuple if isinstance(ele.value, cst.Tuple
) else cst.List,
)) for ele in elements])
else:
# Unrecoginized form
return
self.report(
node,
message_formatter.format(func=call_name),
replacement=node.deep_replace(node, new_node),
)
def obf_universal(self, node: cst.CSTNode, *types):
if m.matches(node, m.Name()):
types = ('a', 'ca', 'v', 'cv') if not types else types
node = cst.ensure_type(node, cst.Name)
if self.can_rename(node.value, *types):
node = self.get_new_cst_name(node)
elif m.matches(node, m.NameItem()):
node = cst.ensure_type(node, cst.NameItem)
node = node.with_changes(name=self.obf_universal(node.name))
elif m.matches(node, m.Call()):
node = cst.ensure_type(node, cst.Call)
if self.change_methods or self.change_functions:
node = self.new_obf_function_name(node)
if self.change_arguments or self.change_method_arguments:
node = self.obf_function_args(node)
elif m.matches(node, m.Attribute()):
node = cst.ensure_type(node, cst.Attribute)
value = node.value
attr = node.attr
self.obf_universal(value)
self.obf_universal(attr)
elif m.matches(node, m.AssignTarget()):
node = cst.ensure_type(node, cst.AssignTarget)
node = node.with_changes(target=self.obf_universal(node.target))
elif m.matches(node, m.List() | m.Tuple()):
node = cst.ensure_type(node, cst.List) if m.matches(
node, m.List()) else cst.ensure_type(node, cst.Tuple)
new_elements = []
for el in node.elements:
new_elements.append(self.obf_universal(el))
node = node.with_changes(elements=new_elements)
elif m.matches(node, m.Subscript()):
node = cst.ensure_type(node, cst.Subscript)
new_slice = []
for el in node.slice:
new_slice.append(
el.with_changes(slice=self.obf_slice(el.slice)))
node = node.with_changes(slice=new_slice)
node = node.with_changes(value=self.obf_universal(node.value))
elif m.matches(node, m.Element()):
node = cst.ensure_type(node, cst.Element)
node = node.with_changes(value=self.obf_universal(node.value))
elif m.matches(node, m.Dict()):
node = cst.ensure_type(node, cst.Dict)
new_elements = []
for el in node.elements:
new_elements.append(self.obf_universal(el))
node = node.with_changes(elements=new_elements)
elif m.matches(node, m.DictElement()):
node = cst.ensure_type(node, cst.DictElement)
new_key = self.obf_universal(node.key)
new_val = self.obf_universal(node.value)
node = node.with_changes(key=new_key, value=new_val)
elif m.matches(node, m.StarredDictElement()):
node = cst.ensure_type(node, cst.StarredDictElement)
node = node.with_changes(value=self.obf_universal(node.value))
elif m.matches(node, m.If() | m.While()):
node = cst.ensure_type(node, cst.IfExp) if m.matches(
node, cst.If
| cst.IfExp) else cst.ensure_type(node, cst.While)
node = node.with_changes(test=self.obf_universal(node.test))
elif m.matches(node, m.IfExp()):
node = cst.ensure_type(node, cst.IfExp)
node = node.with_changes(body=self.obf_universal(node.body))
node = node.with_changes(test=self.obf_universal(node.test))
node = node.with_changes(orelse=self.obf_universal(node.orelse))
elif m.matches(node, m.Comparison()):
node = cst.ensure_type(node, cst.Comparison)
new_compars = []
for target in node.comparisons:
new_compars.append(self.obf_universal(target))
node = node.with_changes(left=self.obf_universal(node.left))
node = node.with_changes(comparisons=new_compars)
elif m.matches(node, m.ComparisonTarget()):
node = cst.ensure_type(node, cst.ComparisonTarget)
node = node.with_changes(
comparator=self.obf_universal(node.comparator))
elif m.matches(node, m.FormattedString()):
node = cst.ensure_type(node, cst.FormattedString)
new_parts = []
for part in node.parts:
new_parts.append(self.obf_universal(part))
node = node.with_changes(parts=new_parts)
elif m.matches(node, m.FormattedStringExpression()):
node = cst.ensure_type(node, cst.FormattedStringExpression)
node = node.with_changes(
expression=self.obf_universal(node.expression))
elif m.matches(node, m.BinaryOperation() | m.BooleanOperation()):
node = cst.ensure_type(node, cst.BinaryOperation) if m.matches(
node, m.BinaryOperation()) else cst.ensure_type(
node, cst.BooleanOperation)
node = node.with_changes(left=self.obf_universal(node.left),
right=self.obf_universal(node.right))
elif m.matches(node, m.UnaryOperation()):
node = cst.ensure_type(node, cst.UnaryOperation)
node = node.with_changes(
expression=self.obf_universal(node.expression))
elif m.matches(node, m.ListComp()):
node = cst.ensure_type(node, cst.ListComp)
node = node.with_changes(elt=self.obf_universal(node.elt))
node = node.with_changes(for_in=self.obf_universal(node.for_in))
elif m.matches(node, m.DictComp()):
node = cst.ensure_type(node, cst.DictComp)
node = node.with_changes(key=self.obf_universal(node.key))
node = node.with_changes(value=self.obf_universal(node.value))
node = node.with_changes(for_in=self.obf_universal(node.for_in))
elif m.matches(node, m.CompFor()):
node = cst.ensure_type(node, cst.CompFor)
new_ifs = []
node = node.with_changes(target=self.obf_universal(node.target))
node = node.with_changes(iter=self.obf_universal(node.iter))
for el in node.ifs:
new_ifs.append(self.obf_universal(el))
node = node.with_changes(ifs=new_ifs)
elif m.matches(node, m.CompIf()):
node = cst.ensure_type(node, cst.CompIf)
node = node.with_changes(test=self.obf_universal(node.test))
elif m.matches(node, m.Integer() | m.Float() | m.SimpleString()):
pass
else:
pass
# print(node)
return node