From 5c6b1efb35c813017b0034742428a3ae5e9128ec Mon Sep 17 00:00:00 2001
From: Michael Milton <michael.r.milton@gmail.com>
Date: Sun, 15 Sep 2024 01:19:03 +1000
Subject: [PATCH 1/2] Add force_class option

---
 src/sqlacodegen/generators.py       |  7 ++++---
 tests/test_generator_declarative.py | 30 +++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py
index 21eadb6..a2a6971 100644
--- a/src/sqlacodegen/generators.py
+++ b/src/sqlacodegen/generators.py
@@ -713,6 +713,7 @@ class DeclarativeGenerator(TablesGenerator):
         "use_inflect",
         "nojoined",
         "nobidi",
+        "force_class"
     }
 
     def __init__(
@@ -774,9 +775,7 @@ def generate_models(self) -> list[Model]:
 
             # Only form model classes for tables that have a primary key and are not
             # association tables
-            if not table.primary_key:
-                models_by_table_name[qualified_name] = Model(table)
-            else:
+            if table.primary_key or "force_class" in self.options:
                 model = ModelClass(table)
                 models_by_table_name[qualified_name] = model
 
@@ -784,6 +783,8 @@ def generate_models(self) -> list[Model]:
                 for column in table.c:
                     column_attr = ColumnAttribute(model, column)
                     model.columns.append(column_attr)
+            else:
+                models_by_table_name[qualified_name] = Model(table)
 
         # Add relationships
         for model in models_by_table_name.values():
diff --git a/tests/test_generator_declarative.py b/tests/test_generator_declarative.py
index d9bf7b5..dcf4e47 100644
--- a/tests/test_generator_declarative.py
+++ b/tests/test_generator_declarative.py
@@ -1509,3 +1509,33 @@ class Simple(Base):
 server_default=text("'test'"))
 """,
     )
+
+
+def test_force_class(
+    metadata: MetaData, engine: Engine
+) -> None:
+    generator = DeclarativeGenerator(metadata, engine, options = ["force_class"])
+    Table(
+        "simple",
+        generator.metadata,
+        Column("a", INTEGER),
+        Column("b", VARCHAR),
+    )
+
+
+    assert generator.generate() == """\
+from typing import Optional
+
+from sqlalchemy import Integer, String
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Simple(Base):
+    __tablename__ = 'simple'
+
+    a: Mapped[Optional[int]] = mapped_column(Integer)
+    b: Mapped[Optional[str]] = mapped_column(String)
+"""

From ee23965f6d11bab472a99de3d54edda734393e5c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 14 Sep 2024 15:26:20 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/sqlacodegen/generators.py       |  2 +-
 tests/test_generator_declarative.py | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py
index a2a6971..56c31c5 100644
--- a/src/sqlacodegen/generators.py
+++ b/src/sqlacodegen/generators.py
@@ -713,7 +713,7 @@ class DeclarativeGenerator(TablesGenerator):
         "use_inflect",
         "nojoined",
         "nobidi",
-        "force_class"
+        "force_class",
     }
 
     def __init__(
diff --git a/tests/test_generator_declarative.py b/tests/test_generator_declarative.py
index dcf4e47..eadcbe5 100644
--- a/tests/test_generator_declarative.py
+++ b/tests/test_generator_declarative.py
@@ -1511,10 +1511,8 @@ class Simple(Base):
     )
 
 
-def test_force_class(
-    metadata: MetaData, engine: Engine
-) -> None:
-    generator = DeclarativeGenerator(metadata, engine, options = ["force_class"])
+def test_force_class(metadata: MetaData, engine: Engine) -> None:
+    generator = DeclarativeGenerator(metadata, engine, options=["force_class"])
     Table(
         "simple",
         generator.metadata,
@@ -1522,8 +1520,9 @@ def test_force_class(
         Column("b", VARCHAR),
     )
 
-
-    assert generator.generate() == """\
+    assert (
+        generator.generate()
+        == """\
 from typing import Optional
 
 from sqlalchemy import Integer, String
@@ -1539,3 +1538,4 @@ class Simple(Base):
     a: Mapped[Optional[int]] = mapped_column(Integer)
     b: Mapped[Optional[str]] = mapped_column(String)
 """
+    )