Add multi-file model output

Author: John Lyu

CHANGES.rst

@@ -5,6 +5,8 @@ Version history
 
 - Added autoincrement to primary key columns to prevent missing field errors.
   (`#473 <https://github.com/agronholm/sqlacodegen/issues/473>`_; PR by @jtmonroe)
+- Added ``--output-directory`` for writing generated models into separate files.
+  (`#88 <https://github.com/agronholm/sqlacodegen/issues/88>`_)
 
 **4.0.3**
 

README.rst

@@ -70,10 +70,15 @@ Examples::
     sqlacodegen postgresql:///some_local_db
     sqlacodegen --generator tables mysql+pymysql://user:password@localhost/dbname
     sqlacodegen --generator dataclasses sqlite:///database.db
+    sqlacodegen sqlite:///database.db --output-directory models
     # --engine-arg values are parsed with ast.literal_eval
     sqlacodegen oracle+oracledb://user:pass@127.0.0.1:1521/XE --engine-arg thick_mode=True
     sqlacodegen oracle+oracledb://user:pass@127.0.0.1:1521/XE --engine-arg thick_mode=True --engine-arg connect_args='{"user": "user", "dsn": "..."}'
 
+To write generated models into separate files, use ``--output-directory``. This creates
+the directory if needed and writes one Python file per generated model, plus shared
+support files such as ``__base.py`` and ``__init__.py``.
+
 To see the list of generic options::
 
     sqlacodegen --help

src/sqlacodegen/cli.py

@@ -5,6 +5,7 @@
 import sys
 from contextlib import ExitStack
 from importlib.metadata import entry_points, version
+from pathlib import Path
 from typing import Any, TextIO
 
 from sqlalchemy.engine import create_engine
@@ -88,7 +89,14 @@ def main() -> None:
             "(values are parsed with ast.literal_eval)"
         ),
     )
-    parser.add_argument("--outfile", help="file to write output to (default: stdout)")
+    output_group = parser.add_mutually_exclusive_group()
+    output_group.add_argument(
+        "--outfile", help="file to write output to (default: stdout)"
+    )
+    output_group.add_argument(
+        "--output-directory",
+        help="directory to write generated models to, one file per model",
+    )
     args = parser.parse_args()
 
     if args.version:
@@ -133,6 +141,14 @@ def main() -> None:
             engine, schema, (generator.views_supported and not args.noviews), tables
         )
 
+    if args.output_directory:
+        output_directory = Path(args.output_directory)
+        output_directory.mkdir(parents=True, exist_ok=True)
+        for name, contents in generator.generate(multi_file=True).items():
+            (output_directory / f"{name}.py").write_text(contents, encoding="utf-8")
+
+        return
+
     # Open the target file (if given)
     with ExitStack() as stack:
         outfile: TextIO

src/sqlacodegen/generators.py

@@ -14,7 +14,7 @@
 from keyword import iskeyword
 from pprint import pformat
 from textwrap import indent
-from typing import Any, ClassVar, Literal, cast
+from typing import Any, ClassVar, Literal, cast, overload
 
 import inflect
 import sqlalchemy
@@ -110,8 +110,16 @@ def __init__(
     def views_supported(self) -> bool:
         pass
 
+    @overload
     @abstractmethod
-    def generate(self) -> str:
+    def generate(self, multi_file: Literal[False] = False) -> str: ...
+
+    @overload
+    @abstractmethod
+    def generate(self, multi_file: Literal[True]) -> dict[str, str]: ...
+
+    @abstractmethod
+    def generate(self, multi_file: bool = False) -> str | dict[str, str]:
         """
         Generate the code for the given metadata.
         .. note:: May modify the metadata.
@@ -167,10 +175,14 @@ def generate_base(self) -> None:
             metadata_ref="metadata",
         )
 
-    def generate(self) -> str:
-        self.generate_base()
+    @overload
+    def generate(self, multi_file: Literal[False] = False) -> str: ...
 
-        sections: list[str] = []
+    @overload
+    def generate(self, multi_file: Literal[True]) -> dict[str, str]: ...
+
+    def generate(self, multi_file: bool = False) -> str | dict[str, str]:
+        self.generate_base()
 
         # Remove unwanted elements from the metadata
         for table in list(self.metadata.tables.values()):
@@ -199,6 +211,14 @@ def generate(self) -> str:
         # Generate the models
         models: list[Model] = self.generate_models()
 
+        if multi_file:
+            return self.render_multi_file(models)
+        else:
+            return self.render_all(models)
+
+    def render_all(self, models: list[Model]) -> str:
+        sections: list[str] = []
+
         # Render module level variables
         if variables := self.render_module_variables(models):
             sections.append(variables + "\n")
@@ -220,6 +240,74 @@ def generate(self) -> str:
 
         return "\n\n".join(sections) + "\n"
 
+    def render_multi_file(self, models: list[Model]) -> dict[str, str]:
+        rendered_files = {
+            "__base": self.render_base_model(),
+            "__init__": "",
+        }
+        base = self.base
+        imports = self.imports
+        module_imports = self.module_imports
+
+        try:
+            for model in models:
+                self.imports = defaultdict(set)
+                self.module_imports = set()
+                self.base = self.get_base_model_import(base, model)
+                self.collect_imports([model])
+                rendered_files[model.name] = self.render_all([model])
+        finally:
+            self.base = base
+            self.imports = imports
+            self.module_imports = module_imports
+
+        return rendered_files
+
+    def render_base_model(self) -> str:
+        imports = self.imports
+        module_imports = self.module_imports
+        try:
+            self.imports = defaultdict(set)
+            self.module_imports = set()
+            for literal_import in self.base.literal_imports:
+                self.add_literal_import(literal_import.pkgname, literal_import.name)
+
+            declarations = list(self.base.declarations)
+            if self.base.table_metadata_declaration is not None:
+                declarations.append(self.base.table_metadata_declaration)
+
+            sections = []
+            groups = self.group_imports()
+            if rendered_imports := "\n\n".join(
+                "\n".join(line for line in group) for group in groups
+            ):
+                sections.append(rendered_imports)
+
+            if declarations:
+                sections.append("\n".join(declarations))
+
+            return "\n\n".join(sections) + "\n"
+        finally:
+            self.imports = imports
+            self.module_imports = module_imports
+
+    def get_base_model_import(self, base: Base, model: Model) -> Base:
+        literal_imports = []
+        if base.metadata_ref == "metadata":
+            literal_imports.append(LiteralImport(".__base", "metadata"))
+        elif base.declarations and isinstance(model, ModelClass):
+            literal_imports.append(
+                LiteralImport(".__base", base.metadata_ref.partition(".")[0])
+            )
+
+        return Base(
+            literal_imports=literal_imports,
+            declarations=[],
+            metadata_ref=base.metadata_ref,
+            decorator=base.decorator,
+            table_metadata_declaration=None,
+        )
+
     def collect_imports(self, models: Iterable[Model]) -> None:
         for literal_import in self.base.literal_imports:
             self.add_literal_import(literal_import.pkgname, literal_import.name)

tests/test_cli.py

@@ -87,6 +87,51 @@ class Foo(Base):
     )
 
 
+def test_cli_declarative_output_directory(db_path: Path, tmp_path: Path) -> None:
+    output_path = tmp_path / "models"
+    subprocess.run(
+        [
+            "sqlacodegen",
+            f"sqlite:///{db_path}",
+            "--generator",
+            "declarative",
+            "--output-directory",
+            str(output_path),
+        ],
+        check=True,
+    )
+
+    assert sorted(path.name for path in output_path.iterdir()) == [
+        "Foo.py",
+        "__base.py",
+        "__init__.py",
+    ]
+    assert (
+        (output_path / "__base.py").read_text()
+        == """\
+from sqlalchemy.orm import DeclarativeBase
+
+class Base(DeclarativeBase):
+    pass
+"""
+    )
+    assert (output_path / "__init__.py").read_text() == ""
+    assert (
+        (output_path / "Foo.py").read_text()
+        == """\
+from .__base import Base
+from sqlalchemy import Integer, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+class Foo(Base):
+    __tablename__ = 'foo'
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    name: Mapped[str] = mapped_column(Text, nullable=False)
+"""
+    )
+
+
 def test_cli_dataclass(db_path: Path, tmp_path: Path) -> None:
     output_path = tmp_path / "outfile"
     subprocess.run(

tests/test_generator_declarative.py

@@ -77,6 +77,54 @@ class SimpleItems(Base):
     )
 
 
+def test_multi_file_generation(metadata: MetaData, engine: Engine) -> None:
+    generator = DeclarativeGenerator(metadata, engine, [])
+    Table(
+        "foo",
+        metadata,
+        Column("id", INTEGER, primary_key=True),
+        Column("name", Text, nullable=False),
+    )
+    Table(
+        "bar",
+        metadata,
+        Column("id", INTEGER, primary_key=True),
+        Column("enabled", INTEGER, nullable=False),
+    )
+
+    assert generator.generate(multi_file=True) == {
+        "__base": """\
+from sqlalchemy.orm import DeclarativeBase
+
+class Base(DeclarativeBase):
+    pass
+""",
+        "__init__": "",
+        "Foo": """\
+from .__base import Base
+from sqlalchemy import Integer, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+class Foo(Base):
+    __tablename__ = 'foo'
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    name: Mapped[str] = mapped_column(Text, nullable=False)
+""",
+        "Bar": """\
+from .__base import Base
+from sqlalchemy import Integer
+from sqlalchemy.orm import Mapped, mapped_column
+
+class Bar(Base):
+    __tablename__ = 'bar'
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    enabled: Mapped[int] = mapped_column(Integer, nullable=False)
+""",
+    }
+
+
 def test_index_with_kwargs(generator: CodeGenerator) -> None:
     simple_items = Table(
         "simple_items",