Use ruff to format generated code

pyproject.toml

@@ -48,6 +48,7 @@ cli = [
     "docformatter>=1.7.2",
     "jinja2>=2.10",
     "toposort>=1.5",
+    "ruff>=0.1.9"
 ]
 docs = [
     "furo",

tests/codegen/test_writer.py

@@ -23,11 +23,15 @@ def setUp(self):
         generator = NoneGenerator(config)
         self.writer = CodeWriter(generator)
 
+    @mock.patch.object(CodeWriter, "ruff_code")
     @mock.patch.object(NoneGenerator, "render_header")
     @mock.patch.object(NoneGenerator, "render")
     @mock.patch.object(NoneGenerator, "normalize_packages")
-    def test_write(self, mock_normalize_packages, mock_render, mock_render_header):
+    def test_write(
+        self, mock_normalize_packages, mock_render, mock_render_header, mock_ruff_code
+    ):
         classes = ClassFactory.list(2)
+        mock_ruff_code.side_effect = lambda x, y: x
         with TemporaryDirectory() as tmpdir:
             mock_render.return_value = [
                 GeneratorResult(Path(f"{tmpdir}/foo/a.py"), "file", "aAa"),
@@ -42,21 +46,33 @@ def test_write(self, mock_normalize_packages, mock_render, mock_render_header):
             self.assertFalse(Path(f"{tmpdir}/c.py").exists())
             mock_normalize_packages.assert_called_once_with(classes)
 
+    @mock.patch.object(CodeWriter, "ruff_code")
+    @mock.patch.object(NoneGenerator, "render_header")
     @mock.patch.object(NoneGenerator, "render")
     @mock.patch.object(NoneGenerator, "normalize_packages")
     @mock.patch("builtins.print")
-    def test_print(self, mock_print, mock_normalize_packages, mock_render):
+    def test_print(
+        self,
+        mock_print,
+        mock_normalize_packages,
+        mock_render,
+        mock_render_header,
+        mock_ruff_code,
+    ):
         classes = ClassFactory.list(2)
         mock_render.return_value = [
             GeneratorResult(Path("foo/a.py"), "file", "aAa"),
             GeneratorResult(Path("bar/b.py"), "file", "bBb"),
             GeneratorResult(Path("c.py"), "file", ""),
         ]
+        mock_render_header.return_value = "# H\n"
+        mock_ruff_code.side_effect = lambda x, y: x
+
         self.writer.print(classes)
 
         mock_normalize_packages.assert_called_once_with(classes)
         mock_print.assert_has_calls(
-            [mock.call("aAa", end=""), mock.call("bBb", end="")]
+            [mock.call("# H\naAa", end=""), mock.call("# H\nbBb", end="")]
         )
 
     def test_from_config(self):
@@ -71,3 +87,42 @@ def test_from_config(self):
         CodeWriter.register_generator("dataclasses", DataclassGenerator)
         writer = CodeWriter.from_config(config)
         self.assertIsInstance(writer.generator, DataclassGenerator)
+
+    def test_ruff_code(self):
+        src_code = (
+            "\n"
+            "import sys\n"
+            "@dataclass\n"
+            "\n"
+            "class MyType:\n"
+            "\n"
+            '    value: Optional[str] = field(default=None, metadata={"type": "Element", "required": True})\n'
+            "\n"
+            "\n"
+            "        "
+        )
+
+        self.writer.generator.config.output.max_line_length = 55
+        actual = self.writer.ruff_code(src_code, Path(__file__))
+        expected = (
+            "import sys\n"
+            "\n"
+            "\n"
+            "@dataclass\n"
+            "class MyType:\n"
+            "    value: Optional[str] = field(\n"
+            "        default=None,\n"
+            '        metadata={"type": "Element", "required": True},\n'
+            "    )\n"
+        )
+        self.assertEqual(expected, actual)
+
+    def test_format_with_invalid_code(self):
+        src_code = """a = "1"""
+        file_path = Path(__file__)
+
+        self.writer.generator.config.output.max_line_length = 55
+        with self.assertRaises(CodeGenerationError) as cm:
+            self.writer.ruff_code(src_code, file_path)
+
+        self.assertIn("Ruff failed", str(cm.exception))

tests/fixtures/annotations/model.py

@@ -1,23 +1,24 @@
 from __future__ import annotations
 from dataclasses import dataclass, field
+from typing import Optional
 from tests.fixtures.annotations.units import unit
 
 __NAMESPACE__ = "http://domain.org/schema/model"
 
 
 @dataclass
 class Measurement:
-    value: float | None = field(
+    value: Optional[float] = field(
         default=None,
         metadata={
             "required": True,
-        }
+        },
     )
-    unit: unit | None = field(
+    unit: Optional[unit] = field(
         default=None,
         metadata={
             "type": "Attribute",
-        }
+        },
     )
 
 

tests/fixtures/annotations/units.py

@@ -5,7 +5,7 @@
 
 
 class unit(Enum):
-    M = 'm'
-    KG = 'kg'
-    VALUE = '%'
-    NA = 'NA'
+    M = "m"
+    KG = "kg"
+    VALUE = "%"
+    NA = "NA"