diff --git a/src/codemodder/codemods/utils_mixin.py b/src/codemodder/codemods/utils_mixin.py index 3c80ca59..42d1a153 100644 --- a/src/codemodder/codemods/utils_mixin.py +++ b/src/codemodder/codemods/utils_mixin.py @@ -35,6 +35,9 @@ def _find_imported_name(self, node: cst.Name): alias.evaluated_name, ): return self.base_name_for_import(import_node, alias) + case BuiltinAssignment(): + return "builtins." + node.value + return node.value def find_base_name(self, node): diff --git a/src/codemodder/utils/utils.py b/src/codemodder/utils/utils.py index 7c8732d1..4b936c2f 100644 --- a/src/codemodder/utils/utils.py +++ b/src/codemodder/utils/utils.py @@ -14,8 +14,6 @@ def list_subclasses(base_kls) -> set[str]: def full_qualified_name_from_class(cls) -> str: - if cls.__module__ == "builtins": - return cls.__qualname__ return f"{cls.__module__}.{cls.__qualname__}" diff --git a/src/core_codemods/remove_debug_breakpoint.py b/src/core_codemods/remove_debug_breakpoint.py index 104dcb07..7a35846f 100644 --- a/src/core_codemods/remove_debug_breakpoint.py +++ b/src/core_codemods/remove_debug_breakpoint.py @@ -21,9 +21,7 @@ def leave_Expr( match call_node := original_node.value: case cst.Call(): - if self.find_base_name( - call_node - ) == "breakpoint" and self.is_builtin_function(call_node): + if self.find_base_name(call_node) == "builtins.breakpoint": self.report_change(original_node) return cst.RemovalSentinel.REMOVE if self.find_base_name(call_node) == "pdb.set_trace": diff --git a/tests/test_nameresolution_mixin.py b/tests/test_nameresolution_mixin.py index 9208b00b..1bf6ce35 100644 --- a/tests/test_nameresolution_mixin.py +++ b/tests/test_nameresolution_mixin.py @@ -84,7 +84,7 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module: tree = cst.parse_module(input_code) TestCodemod(CodemodContext()).transform_module(tree) - def test_get_base_name_no_import(self): + def test_get_base_name_no_assignment(self): class TestCodemod(Codemod, NameResolutionMixin): def transform_module_impl(self, tree: cst.Module) -> cst.Module: stmt = cst.ensure_type(tree.body[-1], cst.SimpleStatementLine) @@ -92,12 +92,31 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module: node = expr.value maybe_name = self.find_base_name(node.func) - assert maybe_name == "exec.capitalize" + assert maybe_name == "foo" return tree input_code = dedent( """\ - exec.capitalize() + foo('hello world') + """ + ) + tree = cst.parse_module(input_code) + TestCodemod(CodemodContext()).transform_module(tree) + + def test_get_base_name_built_in(self): + class TestCodemod(Codemod, NameResolutionMixin): + def transform_module_impl(self, tree: cst.Module) -> cst.Module: + stmt = cst.ensure_type(tree.body[-1], cst.SimpleStatementLine) + expr = cst.ensure_type(stmt.body[0], cst.Expr) + node = expr.value + + maybe_name = self.find_base_name(node.func) + assert maybe_name == "builtins.print" + return tree + + input_code = dedent( + """\ + print('hello world') """ ) tree = cst.parse_module(input_code)