diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index 4cdedc6f..4eafa385 100644 --- a/src/sqlacodegen/generators.py +++ b/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}" diff --git a/src/sqlacodegen/utils.py b/src/sqlacodegen/utils.py index d3844cec..c158bf61 100644 --- a/src/sqlacodegen/utils.py +++ b/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}))