add table comment to generated code

src/sqlacodegen/generators.py

@@ -13,7 +13,7 @@
 from keyword import iskeyword
 from pprint import pformat
 from textwrap import indent
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Dict
 
 import inflect
 import sqlalchemy
@@ -89,6 +89,7 @@ def __init__(
         invalid_options = {opt for opt in options if opt not in self.valid_options}
         if invalid_options:
             raise ValueError("Unrecognized options: " + ", ".join(invalid_options))
+        self.target_version = 1
 
     @abstractmethod
     def generate(self) -> str:
@@ -117,7 +118,7 @@ def __init__(
         self.indentation: str = indentation
         self.imports: dict[str, set[str]] = defaultdict(set)
 
-    def generate(self) -> str:
+    def generate(self, multi_file=False) -> str | Dict[str, str]:
         sections: list[str] = []
 
         # Remove unwanted elements from the metadata
@@ -147,6 +148,37 @@ def generate(self) -> str:
         # Generate the models
         models: list[Model] = self.generate_models()
 
+        if multi_file:
+            sections_dict = {}
+            # create base_model.py and __init__.py
+            sections_dict["base_model"] = """import sqlalchemy as sa
+
+if sa.__version__.split('.')[0] == '1':
+    from sqlalchemy.orm import declarative_base # type: ignore
+else:
+    from sqlalchemy.orm import DeclarativeBase, MappedAsDataclass
+    class DataBase(MappedAsDataclass, DeclarativeBase):
+        pass
+
+    def declarative_base():
+        return DataBase"""
+            sections_dict["__init__"] = f""
+
+            for model in models:
+                # clean imports
+                self.imports: dict[str, set[str]] = defaultdict(set)
+                self.collect_imports([model])
+                if "sqlalchemy.orm" in self.imports and "declarative_base" in self.imports["sqlalchemy.orm"]:
+                    self.remove_literal_import("sqlalchemy.orm", "declarative_base")
+                self.add_literal_import(".base_model", "declarative_base")
+                sections_dict[model.name] = self.render_all([model])
+            return sections_dict
+        else:
+            return self.render_all(models)
+
+    def render_all(self, models: Iterable[Model]) -> str:
+        sections: list[str] = []
+
         # Render module level variables
         variables = self.render_module_variables(models)
         if variables:
@@ -241,7 +273,7 @@ def add_import(self, obj: Any) -> None:
 
             if type_.__name__ in dialect_pkg.__all__:
                 pkgname = dialect_pkgname
-        elif type_.__name__ in sqlalchemy.__all__:  # type: ignore[attr-defined]
+        elif type_.__name__ in sqlalchemy.__dict__:  # type: ignore[attr-defined]
             pkgname = "sqlalchemy"
         else:
             pkgname = type_.__module__
@@ -351,7 +383,7 @@ def render_index(self, index: Index) -> str:
 
         return render_callable("Index", repr(index.name), *extra_args, kwargs=kwargs)
 
-    def render_column(self, column: Column[Any], show_name: bool) -> str:
+    def render_column(self, column: Column[Any], show_name: bool, col_cls="Column") -> str:
         args = []
         kwargs: dict[str, Any] = {}
         kwarg = []
@@ -436,7 +468,7 @@ def render_column(self, column: Column[Any], show_name: bool) -> str:
         if comment:
             kwargs["comment"] = repr(comment)
 
-        return render_callable("Column", *args, kwargs=kwargs)
+        return render_callable(col_cls, *args, kwargs=kwargs)
 
     def render_column_type(self, coltype: object) -> str:
         args = []
@@ -663,6 +695,7 @@ class DeclarativeGenerator(TablesGenerator):
         "use_inflect",
         "nojoined",
         "nobidi",
+        "column_docstring",
     }
 
     def __init__(
@@ -688,12 +721,25 @@ def collect_imports(self, models: Iterable[Model]) -> None:
                 )
             else:
                 self.add_literal_import("sqlalchemy.orm", "declarative_base")
+                self.add_literal_import("sqlalchemy.orm", "Mapped")
+                if self.target_version > 1:
+                    self.add_literal_import("sqlalchemy.orm", "mapped_column")
 
     def collect_imports_for_model(self, model: Model) -> None:
         super().collect_imports_for_model(model)
         if isinstance(model, ModelClass):
             if model.relationships:
                 self.add_literal_import("sqlalchemy.orm", "relationship")
+                self.add_literal_import("sqlalchemy", "ForeignKey")
+
+    def collect_imports_for_column(self, column: Column[Any]) -> None:
+        super().collect_imports_for_column(column)
+        try:
+            python_type = column.type.python_type
+        except NotImplementedError:
+            pass
+        else:
+            self.add_import(python_type)
 
     def generate_models(self) -> list[Model]:
         models_by_table_name: dict[str, Model] = {}
@@ -1036,7 +1082,9 @@ def render_models(self, models: list[Model]) -> str:
 
     def render_class(self, model: ModelClass) -> str:
         sections: list[str] = []
-
+        comments = self.render_table_comment(model)
+        if comments:
+            sections.append(comments)
         # Render class variables / special declarations
         class_vars: str = self.render_class_variables(model)
         if class_vars:
@@ -1075,6 +1123,10 @@ def render_class_declaration(self, model: ModelClass) -> str:
         )
         return f"class {model.name}({parent_class_name}):"
 
+    def render_table_comment(self, model: ModelClass) -> str:
+        if model.table.comment:
+            return f'"""{model.table.comment}"""'
+
     def render_class_variables(self, model: ModelClass) -> str:
         variables = [f"__tablename__ = {model.table.name!r}"]
 
@@ -1129,10 +1181,44 @@ def render_table_args(self, table: Table) -> str:
         else:
             return ""
 
+    def get_column_python_type(self, column: Column[Any]) -> str:
+        try:
+            python_type = column.type.python_type
+        except NotImplementedError:
+            python_type_name = "Any"
+        else:
+            python_type_name = python_type.__name__
+
+        if column.nullable:
+            self.add_literal_import("typing", "Optional")
+            python_type_name = f"Optional[{python_type_name}]"
+        return python_type_name
+
     def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
         column = column_attr.column
-        rendered_column = self.render_column(column, column_attr.name != column.name)
-        return f"{column_attr.name} = {rendered_column}"
+        python_type_name = self.get_column_python_type(column)
+        colargs = {}
+        if self.target_version > 1:
+            colargs.update({"col_cls" : "mapped_column"})
+        rendered_column = self.render_column(column, column_attr.name != column.name, **colargs)
+        # kwargs["metadata"] = f"{{{self.metadata_key!r}: {rendered_column}}}"
+        # rendered_field = render_callable("field", kwargs=kwargs)
+        col_line = f"{column_attr.name}: Mapped[{python_type_name}] = {rendered_column}"
+        if "column_docstring" in self.options:
+            docstring = self.render_column_docstring(column)
+            if docstring:
+                return f"{col_line}\n{docstring}"
+        return col_line
+
+    def render_column_docstring(self, column: Column[Any]) -> str:
+        if isinstance(column.comment, str):
+            column.comment = column.comment.replace("\\", "/")
+            comment_line = repr(column.comment)
+            return comment_line
+        else:
+            return ""
+
+
 
     def render_relationship(self, relationship: RelationshipAttribute) -> str:
         def render_column_attrs(column_attrs: list[ColumnAttribute]) -> str:
@@ -1260,15 +1346,6 @@ def collect_imports_for_model(self, model: Model) -> None:
                 ):
                     self.add_literal_import("typing", "List")
 
-    def collect_imports_for_column(self, column: Column[Any]) -> None:
-        super().collect_imports_for_column(column)
-        try:
-            python_type = column.type.python_type
-        except NotImplementedError:
-            pass
-        else:
-            self.add_import(python_type)
-
     def render_module_variables(self, models: list[Model]) -> str:
         if not any(isinstance(model, ModelClass) for model in models):
             return super().render_module_variables(models)
@@ -1311,8 +1388,12 @@ def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
             self.add_literal_import("typing", "Optional")
             kwargs["default"] = None
             python_type_name = f"Optional[{python_type_name}]"
-
-        rendered_column = self.render_column(column, column_attr.name != column.name)
+
+        if self.target_version > 1:
+            colargs = {"col_cls" : "mapped_column"}
+        else:
+            colargs = {}
+        rendered_column = self.render_column(column, column_attr.name != column.name, **colargs)
         kwargs["metadata"] = f"{{{self.metadata_key!r}: {rendered_column}}}"
         rendered_field = render_callable("field", kwargs=kwargs)
         return f"{column_attr.name}: {python_type_name} = {rendered_field}"

src/sqlacodegen/utils.py

@@ -4,7 +4,7 @@
 from collections.abc import Mapping
 
 from sqlalchemy import PrimaryKeyConstraint, UniqueConstraint
-from sqlalchemy.engine import Connectable
+from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.sql.elements import TextClause
 from sqlalchemy.sql.schema import (
@@ -33,7 +33,7 @@ def get_constraint_sort_key(constraint: Constraint) -> str:
         return str(constraint)
 
 
-def get_compiled_expression(statement: ClauseElement, bind: Connectable) -> str:
+def get_compiled_expression(statement: ClauseElement, bind: Connection | Engine) -> str:
     """Return the statement in a form where any placeholders have been filled in."""
     return str(statement.compile(bind, compile_kwargs={"literal_binds": True}))