def test_and_matcher_false(self) -> None: # Fail to match since True and False cannot match. self.assertFalse( matches(cst.Name("None"), m.AllOf(m.Name("True"), m.Name("False"))) ) self.assertFalse( 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.AllOf( (m.Arg(), m.Arg(), m.Arg()), ( m.Arg(m.Integer("3")), m.Arg(m.Integer("2")), m.Arg(m.Integer("1")), ), ), ), ) )
def test_and_matcher_true(self) -> None: # Match on True identifier in roundabout way. self.assertTrue( matches( cst.Name("True"), m.AllOf(m.Name(), m.Name(value=m.MatchRegex(r"True"))) ) ) 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.AllOf( (m.Arg(), m.Arg(), m.Arg()), ( m.Arg(m.Integer("1")), m.Arg(m.Integer("2")), m.Arg(m.Integer("3")), ), ), ), ) )
def test_at_least_n_matcher_args_false(self) -> None: # Fail to match a function call to "foo" where the first argument is the # integer value 1, and there are at least two arguments after that are # strings. self.assertFalse( 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.SimpleString()), n=2), ), ), )) # Fail to match a function call to "foo" where the first argument is the integer # value 1, and there are at least three wildcard arguments after. self.assertFalse( 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=3)), ), )) # Fail to 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 after. self.assertFalse( 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("2")), n=2), ), ), ))
def test_at_most_n_matcher_no_args_true(self) -> None: # Match a function call to "foo" with at most two arguments. self.assertTrue( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Integer("1")), )), m.Call(m.Name("foo"), (m.AtMostN(n=2), )), )) # Match a function call to "foo" with at most two arguments. self.assertTrue( matches( libcst.Call( libcst.Name("foo"), (libcst.Arg( libcst.Integer("1")), libcst.Arg(libcst.Integer("2"))), ), m.Call(m.Name("foo"), (m.AtMostN(n=2), )), )) # Match a function call to "foo" with at most six arguments, the last # one being the integer 1. self.assertTrue( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Integer("1")), )), m.Call(m.Name("foo"), [m.AtMostN(n=5), m.Arg(m.Integer("1"))]), )) # Match a function call to "foo" with at most six arguments, the last # one being the integer 1. self.assertTrue( matches( libcst.Call( libcst.Name("foo"), (libcst.Arg( libcst.Integer("1")), libcst.Arg(libcst.Integer("2"))), ), m.Call(m.Name("foo"), (m.AtMostN(n=5), m.Arg(m.Integer("2")))), )) # Match a function call to "foo" with at most six arguments, the first # one being the integer 1. self.assertTrue( matches( libcst.Call( libcst.Name("foo"), (libcst.Arg( libcst.Integer("1")), libcst.Arg(libcst.Integer("2"))), ), m.Call(m.Name("foo"), (m.Arg(m.Integer("1")), m.AtMostN(n=5))), )) # Match a function call to "foo" with at most six arguments, the first # one being the integer 1. self.assertTrue( matches( libcst.Call( libcst.Name("foo"), (libcst.Arg( libcst.Integer("1")), libcst.Arg(libcst.Integer("2"))), ), m.Call(m.Name("foo"), (m.Arg(m.Integer("1")), m.ZeroOrOne())), ))
def test_at_least_n_matcher_no_args_false(self) -> None: # Fail to match a function call to "foo" with at least four arguments. self.assertFalse( 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.AtLeastN(n=4),)), ) ) # Fail to match a function call to "foo" with at least four arguments, # the first one being the value 1. self.assertFalse( 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(n=3)) ), ) ) # Fail to match a function call to "foo" with at least three arguments, # the last one being the value 2. self.assertFalse( 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.AtLeastN(n=2), m.Arg(m.Integer("2"))) ), ) )
def test_replace_add_one_to_foo_args(self) -> None: def _add_one_to_arg( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return node.deep_replace( # This can be either a node or a sequence, pyre doesn't know. cst.ensure_type(extraction["arg"], cst.CSTNode), # Grab the arg and add one to its value. cst.Integer( str( int( cst.ensure_type(extraction["arg"], cst.Integer).value) + 1)), ) # Verify way more complex transform behavior. original = cst.parse_module( "foo: int = 37\ndef bar(baz: int) -> int:\n return baz\n\nbiz: int = bar(41)\n" ) replaced = cst.ensure_type( m.replace( original, m.Call( func=m.Name("bar"), args=[m.Arg(m.SaveMatchedNode(m.Integer(), "arg"))], ), _add_one_to_arg, ), cst.Module, ).code self.assertEqual( replaced, "foo: int = 37\ndef bar(baz: int) -> int:\n return baz\n\nbiz: int = bar(42)\n", )
def test_predicate_logic_operators_on_attributes(self) -> None: # Verify that we can or things together. matcher = m.BinaryOperation(left=m.Name( metadata=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.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.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.MatchMetadata( meta.ExpressionContextProvider, meta.ExpressionContext.STORE)))) node, wrapper = self._make_fixture("a + b") self.assertTrue(matches(node, matcher, metadata_resolver=wrapper))
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 test_at_most_n_matcher_no_args_false(self) -> None: # Fail to match a function call to "foo" with at most two arguments. self.assertFalse( 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.AtMostN(n=2),)), ) ) # Fail to match a function call to "foo" with at most two arguments, # the last one being the integer 3. self.assertFalse( 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.AtMostN(n=1), m.Arg(m.Integer("3"))) ), ) ) # Fail to match a function call to "foo" with at most two arguments, # the last one being the integer 3. self.assertFalse( 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.ZeroOrOne(), m.Arg(m.Integer("3")))), ) )
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_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_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 transform_module_impl(self, tree: cst.Module) -> cst.Module: if self.iterations == 0: return tree self.iterations -= 1 return cst.ensure_type( m.replace( tree, m.Integer(), lambda node, _: node.with_changes(value=str(int(node.value) + 1)), ), cst.Module, )
def test_zero_or_more_matcher_args_false(self) -> None: # Fail to match a function call to "foo" where the first argument is the # integer value 1, and the rest of the arguments are strings. self.assertFalse( 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.SimpleString()))), ), ) ) # Fail to 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. self.assertFalse( 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("2")))), ), ) )
def test_replace_add_one(self) -> None: def _add_one( node: cst.CSTNode, extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]], ) -> cst.CSTNode: return cst.Integer( str(int(cst.ensure_type(node, cst.Integer).value) + 1)) # Verify slightly more complex transform behavior. original = cst.parse_module( "foo: int = 36\ndef bar() -> int:\n return 41\n") replaced = cst.ensure_type(m.replace(original, m.Integer(), _add_one), cst.Module).code self.assertEqual(replaced, "foo: int = 37\ndef bar() -> int:\n return 42\n")
def test_at_most_n_matcher_args_false(self) -> None: # Fail to match a function call to "foo" with at most three arguments, # all of which are the integer 4. 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.AtMostN(m.Arg(m.Integer("4")), n=3), )), ))
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 test_does_not_match_operator_false(self) -> None: # Match on any call that takes one argument that isn't the value None. self.assertFalse( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Name("None")), )), m.Call(args=(m.Arg(value=~m.Name("None")), )), )) self.assertFalse( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Integer("1")), )), m.Call(args=((~m.Arg(m.Integer("1"))), )), )) # Match any call that takes an argument which isn't True or False. self.assertFalse( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Name("False")), )), m.Call(args=(m.Arg( value=~(m.Name("True") | m.Name("False"))), )), )) # Roundabout way of verifying ~(x&y) behavior. self.assertFalse( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Name("False")), )), m.Call(args=(m.Arg(value=~(m.Name() & m.Name("False"))), )), )) # Roundabout way of verifying (~x)|(~y) behavior self.assertFalse( matches( libcst.Call(libcst.Name("foo"), (libcst.Arg(libcst.Name("True")), )), m.Call(args=(m.Arg(value=(~m.Name("True")) | (~m.Name("True"))), )), )) # Match any name node that doesn't match the regex for True self.assertFalse( matches(libcst.Name("True"), m.Name(value=~m.MatchRegex(r"True"))))
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 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_replace_add_one() -> None: original = cst.parse_module( "foo: int = 36\ndef bar() -> int:\n return 41\n") replaced = m.replace(original, m.Integer(), _add_one) assert replaced == "foo: int = 37\ndef bar() -> int:\n return 42\n"
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 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
def test_at_least_n_matcher_no_args_true(self) -> None: # Match a function call to "foo" with at least one argument. 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.AtLeastN(n=1),)), ) ) # Match a function call to "foo" with at least two arguments. 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.AtLeastN(n=2),)), ) ) # Match a function call to "foo" with at least three arguments. 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.AtLeastN(n=3),)), ) ) # Match a function call to "foo" with at least two arguments the # first one being the integer 1. 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(n=1)) ), ) ) # Match a function call to "foo" with at least three arguments the # first one being the integer 1. 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(n=2)) ), ) ) # Match a function call to "foo" with at least three arguments. The # There should be an argument with the value 2, which should have # at least one argument before and one argument 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.AtLeastN(n=1), m.Arg(m.Integer("2")), m.AtLeastN(n=1)), ), ) ) # Match a function call to "foo" with at least two arguments, the last # one being the value 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.AtLeastN(n=1), m.Arg(m.Integer("3"))) ), ) ) # Match a function call to "foo" with at least three arguments, the last # one being the value 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.AtLeastN(n=2), m.Arg(m.Integer("3"))) ), ) )
def test_zero_or_more_matcher_no_args_true(self) -> None: # Match a function call to "foo" with any number of arguments as # long as the first one is an integer with the value 1. 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()) ), ) ) # Match a function call to "foo" with any number of arguments as # long as one of them is an integer with the value 1. 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.Integer("1")), m.ZeroOrMore()), ), ) ) # Match a function call to "foo" with any number of arguments as # long as one of them is an integer with the value 2. 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.Integer("2")), m.ZeroOrMore()), ), ) ) # Match a function call to "foo" with any number of arguments as # long as one of them is an integer with the value 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.ZeroOrMore(), m.Arg(m.Integer("3")), m.ZeroOrMore()), ), ) ) # Match a function call to "foo" with any number of arguments as # long as the last one is an integer with the value 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.ZeroOrMore(), m.Arg(m.Integer("3"))) ), ) ) # Match a function call to "foo" with any number of arguments as # long as there are two arguments with the values 1 and 3 anywhere # in the argument list, respecting order. 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.Integer("1")), m.ZeroOrMore(), m.Arg(m.Integer("3")), m.ZeroOrMore(), ), ), ) ) # Match a function call to "foo" with any number of arguments as # long as there are three arguments with the values 1, 2 and 3 anywhere # in the argument list, respecting order. 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.Integer("1")), m.ZeroOrMore(), m.Arg(m.Integer("2")), m.ZeroOrMore(), m.Arg(m.Integer("3")), m.ZeroOrMore(), ), ), ) )
def test_complex_matcher_false(self) -> None: # Fail to match since this is a Call, not a FunctionDef. self.assertFalse( matches( cst.Call( func=cst.Name("foo"), args=( cst.Arg(cst.Integer("1")), cst.Arg(cst.Integer("2")), cst.Arg(cst.Integer("3")), ), ), m.FunctionDef(), ) ) # Fail to match a function named "bar". self.assertFalse( 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(m.Name("bar")), ) ) # Fail to match a function named "foo" with two arguments. self.assertFalse( 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.Arg())), ) ) # Fail to match a function named "foo" with three integer arguments # 3, 2, 1. self.assertFalse( 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("3")), m.Arg(m.Integer("2")), m.Arg(m.Integer("1")), ), ), ) ) # Fail to match a function named "foo" with three arguments, the last one # being the integer 1. self.assertFalse( 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.DoNotCare(), m.DoNotCare(), m.Arg(m.Integer("1"))), ), ) )
def remove_trailing_comma(node): # Remove the comma from this node, *unless* it's already a comma node with comments if node.comma is cst.MaybeSentinel.DEFAULT or m.findall(node, m.Comment()): return node return node.with_changes(comma=cst.MaybeSentinel.DEFAULT) MATCH_NONE = m.MatchIfTrue(lambda x: x is None) ALL_ELEMS_SLICE = m.Slice( lower=MATCH_NONE | m.Name("None"), upper=MATCH_NONE | m.Name("None"), step=MATCH_NONE | m.Name("None") | m.Integer("1") | m.UnaryOperation(m.Minus(), m.Integer("1")), ) class ShedFixers(VisitorBasedCodemodCommand): """Fix a variety of small problems. Replaces `raise NotImplemented` with `raise NotImplementedError`, and converts always-failing assert statements to explicit `raise` statements. Also includes code closely modelled on pybetter's fixers, because it's considerably faster to run all transforms in a single pass if possible. """ DESCRIPTION = "Fix a variety of style, performance, and correctness issues."
def visit_Call(self, node: cst.Call) -> None: result = m.extract( node, m.Call( func=m.Attribute(value=m.Name("self"), attr=m.Name("assertTrue")), args=[ m.DoNotCare(), m.Arg(value=m.SaveMatchedNode( m.OneOf( m.Integer(), m.Float(), m.Imaginary(), m.Tuple(), m.List(), m.Set(), m.Dict(), m.Name("None"), m.Name("True"), m.Name("False"), ), "second", )), ], ), ) if result: second_arg = result["second"] if isinstance(second_arg, Sequence): second_arg = second_arg[0] if m.matches(second_arg, m.Name("True")): new_call = node.with_changes(args=[ node.args[0].with_changes(comma=cst.MaybeSentinel.DEFAULT) ], ) elif m.matches(second_arg, m.Name("None")): new_call = node.with_changes( func=node.func.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertIsNone", ), args=[ node.args[0].with_changes( comma=cst.MaybeSentinel.DEFAULT) ], ) elif m.matches(second_arg, m.Name("False")): new_call = node.with_changes( func=node.func.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertFalse", ), args=[ node.args[0].with_changes( comma=cst.MaybeSentinel.DEFAULT) ], ) else: new_call = node.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertEqual", ) self.report(node, replacement=new_call)
def test_complex_matcher_true(self) -> None: # Match on any Call, not caring about arguments. 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(), ) ) # Match on any Call to a function named "foo". 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(m.Name("foo")), ) ) # Match on any Call to a function named "foo" with three arguments. 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.Arg(), m.Arg())), ) ) # Match any Call to a function named "foo" with three integer arguments. 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()), m.Arg(m.Integer()), m.Arg(m.Integer())), ), ) ) # Match any Call to a function named "foo" with integer arguments 1, 2, 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.Arg(m.Integer("2")), m.Arg(m.Integer("3")), ), ), ) ) # Match any Call to a function named "foo" with three arguments, the last one # being the integer 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.DoNotCare(), m.DoNotCare(), m.Arg(m.Integer("3"))), ), ) )
def _deprecated_visitor(self, node: cst.ExtSlice) -> None: if m.matches(node, m.ExtSlice(m.Index(m.Integer("2")))): self.calls.add("called")