fix-dataclass-defaults can handle populated defaults (#479)

Author: clavedeluna

src/codemodder/codemods/base_visitor.py

@@ -66,6 +66,15 @@ def node_position(self, node):
     def lineno_for_node(self, node):
         return self.node_position(node).start.line
 
+    def code(self, node: cst.CSTNode) -> str:
+        """
+        Only a cst.Module node has a `code` attribute which converts the node
+        back to the original code as a str. To get the code for any node,
+        the suggested approach is to wrap this node in a `cst.Module` node.
+        """
+        module = cst.Module(body=[cst.SimpleStatementLine(body=[cst.Expr(value=node)])])
+        return module.code
+
 
 class BaseTransformer(VisitorBasedCodemodCommand, UtilsMixin):
     def __init__(

src/core_codemods/fix_dataclass_defaults.py

@@ -36,23 +36,44 @@ def leave_AnnAssign(
             return updated_node
 
         match original_node.value:
-            # TODO: add support for populated elements
             case cst.List(elements=[]) | cst.Dict(elements=[]) | cst.Tuple(elements=[]):
-                self.add_needed_import("dataclasses", "field")
-                self.report_change(original_node)
-                return updated_node.with_changes(
-                    value=cst.parse_expression(
-                        f"field(default_factory={ type(original_node.value).__name__.lower()})"
-                    )
+                return self.field_with_default_factory(original_node, updated_node)
+            case (
+                cst.List(elements=[_, *_])
+                | cst.Dict(elements=[_, *_])
+                | cst.Tuple(elements=[_, *_])
+            ):
+                return self.field_with_default_factory(
+                    original_node, updated_node, empty=False
                 )
             case cst.Call(func=cst.Name(value="set"), args=[]):
-                self.add_needed_import("dataclasses", "field")
-                self.report_change(original_node)
-                return updated_node.with_changes(
-                    value=cst.parse_expression("field(default_factory=set)")
+                return self.field_with_default_factory(original_node, updated_node)
+            case cst.Call(func=cst.Name(value="set"), args=[_, *_]):
+                return self.field_with_default_factory(
+                    original_node, updated_node, empty=False
                 )
         return updated_node
 
+    def field_with_default_factory(
+        self,
+        original_node: cst.List | cst.Tuple | cst.Dict | cst.Call,
+        updated_node: cst.List | cst.Tuple | cst.Dict | cst.Call,
+        empty=True,
+    ):
+        self.add_needed_import("dataclasses", "field")
+        self.report_change(original_node)
+        value = original_node.value
+        if empty:
+            expr = (
+                "field(default_factory=set)"
+                if isinstance(value, cst.Call)
+                else f"field(default_factory={type(value).__name__.lower()})"
+            )
+            return updated_node.with_changes(value=cst.parse_expression(expr))
+
+        expr = f"field(default_factory=lambda: {self.code(value).strip()})"
+        return updated_node.with_changes(value=cst.parse_expression(expr))
+
     def _has_dataclass_decorator(self, node: cst.ClassDef) -> bool:
         for decorator in node.decorators:
             if self.find_base_name(decorator.decorator) == "dataclasses.dataclass":

tests/codemods/test_base_visitor.py

@@ -55,6 +55,34 @@ def leave_FunctionDef(
         return updated_node
 
 
+class AssertNodeCode(BaseTransformer):
+    def __init__(
+        self,
+        context,
+        results,
+        expected_code,
+        node_type,
+        line_exclude=None,
+        line_include=None,
+    ):
+        BaseTransformer.__init__(
+            self, context, results, line_include or [], line_exclude or []
+        )
+        self.expected_code = expected_code
+        self.node_type = node_type
+        self.checked = False
+
+    def on_leave(
+        self, original_node: cst.CSTNode, updated_node: cst.CSTNode
+    ) -> cst.CSTNode:
+        match original_node:
+            case self.node_type():
+                assert self.code(original_node).strip() == self.expected_code.strip()
+                self.checked = True
+
+        return updated_node
+
+
 class TestBaseVisitor:
     def run_and_assert(self, input_code, expected, line_exclude, line_include):
         input_tree = cst.parse_module(input_code)
@@ -125,3 +153,35 @@ def hello(one, *args, **kwargs):
             start=CodePosition(line=2, column=0), end=CodePosition(line=2, column=31)
         )
         self.run_and_assert(input_code, expected_pos)
+
+
+class TestCodeForNode:
+    def run_and_assert(self, input_code, expected_code, node_type):
+        input_tree = cst.parse_module(dedent(input_code))
+        command_instance = AssertNodeCode(
+            CodemodContext(), defaultdict(list), expected_code, node_type
+        )
+        command_instance.transform_module(input_tree)
+        if not command_instance.checked:
+            raise Exception(f"Input code does not contain a {node_type}")
+
+    def test_annAssign(self):
+        input_code = "phones: list = [1, 2, 3]"
+        self.run_and_assert(input_code, input_code, cst.AnnAssign)
+
+    def test_list(self):
+        input_code = "[1, 'two', Exception]"
+        self.run_and_assert(input_code, input_code, cst.List)
+
+    def test_dict(self):
+        self.run_and_assert('dict = {"friend": "one"}', '{"friend": "one"}', cst.Dict)
+
+    def test_module(self):
+        input_code = dedent(
+            """
+        # comment
+        var = 1
+        print(var)
+        """
+        )
+        self.run_and_assert(input_code, input_code, cst.Module)

tests/codemods/test_fix_dataclass_defaults.py

@@ -58,7 +58,6 @@ class Test:
         self.run_and_assert(tmpdir, input_code, expected, num_changes=3)
 
     def test_populated_defaults(self, tmpdir):
-        # TODO: support later using lambda.
         input_code = """
         import dataclasses
 
@@ -69,7 +68,18 @@ class Test:
             friends: dict = {"friend": "one"}
             family: set = set((1, 2, 3))
         """
-        self.run_and_assert(tmpdir, input_code, input_code)
+        expected = """
+        import dataclasses
+        from dataclasses import field
+        
+        @dataclasses.dataclass
+        class Test:
+            name: str = ""
+            phones: list = field(default_factory=lambda: [1, 2, 3])
+            friends: dict = field(default_factory=lambda: {"friend": "one"})
+            family: set = field(default_factory=lambda: set((1, 2, 3)))
+        """
+        self.run_and_assert(tmpdir, input_code, expected, num_changes=3)
 
     @pytest.mark.parametrize(
         "code",