• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

agronholm / sqlacodegen / 19600858730

22 Nov 2025 08:32PM UTC coverage: 97.365% (-0.3%) from 97.64%
19600858730

Pull #438

github

web-flow
Merge 0206ea8e1 into d7a6024df
Pull Request #438: Add support for rendering dialect kwargs and info, and introduce keep-dialect-types option

74 of 80 new or added lines in 3 files covered. (92.5%)

22 existing lines in 1 file now uncovered.

1515 of 1556 relevant lines covered (97.37%)

4.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.61
/src/sqlacodegen/generators.py
1
from __future__ import annotations
5✔
2

3
import inspect
5✔
4
import re
5✔
5
import sys
5✔
6
from abc import ABCMeta, abstractmethod
5✔
7
from collections import defaultdict
5✔
8
from collections.abc import Collection, Iterable, Sequence
5✔
9
from dataclasses import dataclass
5✔
10
from importlib import import_module
5✔
11
from inspect import Parameter
5✔
12
from itertools import count
5✔
13
from keyword import iskeyword
5✔
14
from pprint import pformat
5✔
15
from textwrap import indent
5✔
16
from typing import Any, ClassVar, Literal, cast
5✔
17

18
import inflect
5✔
19
import sqlalchemy
5✔
20
from sqlalchemy import (
5✔
21
    ARRAY,
22
    Boolean,
23
    CheckConstraint,
24
    Column,
25
    Computed,
26
    Constraint,
27
    DefaultClause,
28
    Enum,
29
    ForeignKey,
30
    ForeignKeyConstraint,
31
    Identity,
32
    Index,
33
    MetaData,
34
    PrimaryKeyConstraint,
35
    String,
36
    Table,
37
    Text,
38
    TypeDecorator,
39
    UniqueConstraint,
40
)
41
from sqlalchemy.dialects.postgresql import DOMAIN, JSON, JSONB
5✔
42
from sqlalchemy.engine import Connection, Engine
5✔
43
from sqlalchemy.exc import CompileError
5✔
44
from sqlalchemy.sql.elements import TextClause
5✔
45
from sqlalchemy.sql.type_api import UserDefinedType
5✔
46
from sqlalchemy.types import TypeEngine
5✔
47

48
from .models import (
5✔
49
    ColumnAttribute,
50
    JoinType,
51
    Model,
52
    ModelClass,
53
    RelationshipAttribute,
54
    RelationshipType,
55
)
56
from .utils import (
5✔
57
    decode_postgresql_sequence,
58
    get_column_names,
59
    get_common_fk_constraints,
60
    get_compiled_expression,
61
    get_constraint_sort_key,
62
    get_stdlib_module_names,
63
    qualified_table_name,
64
    render_callable,
65
    uses_default_name,
66
)
67

68
_re_boolean_check_constraint = re.compile(r"(?:.*?\.)?(.*?) IN \(0, 1\)")
5✔
69
_re_column_name = re.compile(r'(?:(["`]?).*\1\.)?(["`]?)(.*)\2')
5✔
70
_re_enum_check_constraint = re.compile(r"(?:.*?\.)?(.*?) IN \((.+)\)")
5✔
71
_re_enum_item = re.compile(r"'(.*?)(?<!\\)'")
5✔
72
_re_invalid_identifier = re.compile(r"(?u)\W")
5✔
73

74

75
@dataclass
5✔
76
class LiteralImport:
5✔
77
    pkgname: str
5✔
78
    name: str
5✔
79

80

81
@dataclass
5✔
82
class Base:
5✔
83
    """Representation of MetaData for Tables, respectively Base for classes"""
84

85
    literal_imports: list[LiteralImport]
5✔
86
    declarations: list[str]
5✔
87
    metadata_ref: str
5✔
88
    decorator: str | None = None
5✔
89
    table_metadata_declaration: str | None = None
5✔
90

91

92
class CodeGenerator(metaclass=ABCMeta):
5✔
93
    valid_options: ClassVar[set[str]] = set()
5✔
94

95
    def __init__(
5✔
96
        self, metadata: MetaData, bind: Connection | Engine, options: Sequence[str]
97
    ):
98
        self.metadata: MetaData = metadata
5✔
99
        self.bind: Connection | Engine = bind
5✔
100
        self.options: set[str] = set(options)
5✔
101

102
        # Validate options
103
        invalid_options = {opt for opt in options if opt not in self.valid_options}
5✔
104
        if invalid_options:
5✔
105
            raise ValueError("Unrecognized options: " + ", ".join(invalid_options))
×
106

107
    @property
5✔
108
    @abstractmethod
5✔
109
    def views_supported(self) -> bool:
5✔
110
        pass
×
111

112
    @abstractmethod
5✔
113
    def generate(self) -> str:
5✔
114
        """
115
        Generate the code for the given metadata.
116
        .. note:: May modify the metadata.
117
        """
118

119

120
@dataclass(eq=False)
5✔
121
class TablesGenerator(CodeGenerator):
5✔
122
    valid_options: ClassVar[set[str]] = {
5✔
123
        "noindexes",
124
        "noconstraints",
125
        "nocomments",
126
        # accept both spellings
127
        "include_dialect_options",
128
        "include-dialect-options",
129
        "keep_dialect_types",
130
        "keep-dialect-types",
131
    }
132
    stdlib_module_names: ClassVar[set[str]] = get_stdlib_module_names()
5✔
133

134
    def __init__(
5✔
135
        self,
136
        metadata: MetaData,
137
        bind: Connection | Engine,
138
        options: Sequence[str],
139
        *,
140
        indentation: str = "    ",
141
    ):
142
        super().__init__(metadata, bind, options)
5✔
143
        self.indentation: str = indentation
5✔
144
        self.imports: dict[str, set[str]] = defaultdict(set)
5✔
145
        self.module_imports: set[str] = set()
5✔
146

147
        # Render SchemaItem.info and dialect kwargs (Table/Column) into output
148
        self.include_dialect_options_and_info: bool = (
5✔
149
            "include_dialect_options" in self.options
150
            or "include-dialect-options" in self.options
151
        )
152
        # Keep dialect-specific types instead of adapting to generic SQLAlchemy types
153
        self.keep_dialect_types: bool = (
5✔
154
            "keep_dialect_types" in self.options or "keep-dialect-types" in self.options
155
        )
156

157
    @property
5✔
158
    def views_supported(self) -> bool:
5✔
UNCOV
159
        return True
×
160

161
    def generate_base(self) -> None:
5✔
162
        self.base = Base(
5✔
163
            literal_imports=[LiteralImport("sqlalchemy", "MetaData")],
164
            declarations=["metadata = MetaData()"],
165
            metadata_ref="metadata",
166
        )
167

168
    def generate(self) -> str:
5✔
169
        self.generate_base()
5✔
170

171
        sections: list[str] = []
5✔
172

173
        # Remove unwanted elements from the metadata
174
        for table in list(self.metadata.tables.values()):
5✔
175
            if self.should_ignore_table(table):
5✔
UNCOV
176
                self.metadata.remove(table)
×
UNCOV
177
                continue
×
178

179
            if "noindexes" in self.options:
5✔
180
                table.indexes.clear()
5✔
181

182
            if "noconstraints" in self.options:
5✔
183
                table.constraints.clear()
5✔
184

185
            if "nocomments" in self.options:
5✔
186
                table.comment = None
5✔
187

188
            for column in table.columns:
5✔
189
                if "nocomments" in self.options:
5✔
190
                    column.comment = None
5✔
191

192
        # Use information from column constraints to figure out the intended column
193
        # types
194
        for table in self.metadata.tables.values():
5✔
195
            self.fix_column_types(table)
5✔
196

197
        # Generate the models
198
        models: list[Model] = self.generate_models()
5✔
199

200
        # Render module level variables
201
        variables = self.render_module_variables(models)
5✔
202
        if variables:
5✔
203
            sections.append(variables + "\n")
5✔
204

205
        # Render models
206
        rendered_models = self.render_models(models)
5✔
207
        if rendered_models:
5✔
208
            sections.append(rendered_models)
5✔
209

210
        # Render collected imports
211
        groups = self.group_imports()
5✔
212
        imports = "\n\n".join("\n".join(line for line in group) for group in groups)
5✔
213
        if imports:
5✔
214
            sections.insert(0, imports)
5✔
215

216
        return "\n\n".join(sections) + "\n"
5✔
217

218
    def collect_imports(self, models: Iterable[Model]) -> None:
5✔
219
        for literal_import in self.base.literal_imports:
5✔
220
            self.add_literal_import(literal_import.pkgname, literal_import.name)
5✔
221

222
        for model in models:
5✔
223
            self.collect_imports_for_model(model)
5✔
224

225
    def collect_imports_for_model(self, model: Model) -> None:
5✔
226
        if model.__class__ is Model:
5✔
227
            self.add_import(Table)
5✔
228

229
        for column in model.table.c:
5✔
230
            self.collect_imports_for_column(column)
5✔
231

232
        for constraint in model.table.constraints:
5✔
233
            self.collect_imports_for_constraint(constraint)
5✔
234

235
        for index in model.table.indexes:
5✔
236
            self.collect_imports_for_constraint(index)
5✔
237

238
    def collect_imports_for_column(self, column: Column[Any]) -> None:
5✔
239
        self.add_import(column.type)
5✔
240

241
        if isinstance(column.type, ARRAY):
5✔
242
            self.add_import(column.type.item_type.__class__)
5✔
243
        elif isinstance(column.type, (JSONB, JSON)):
5✔
244
            if (
5✔
245
                not isinstance(column.type.astext_type, Text)
246
                or column.type.astext_type.length is not None
247
            ):
248
                self.add_import(column.type.astext_type)
5✔
249
        elif isinstance(column.type, DOMAIN):
5✔
250
            self.add_import(column.type.data_type.__class__)
5✔
251

252
        if column.default:
5✔
253
            self.add_import(column.default)
5✔
254

255
        if column.server_default:
5✔
256
            if isinstance(column.server_default, (Computed, Identity)):
5✔
257
                self.add_import(column.server_default)
5✔
258
            elif isinstance(column.server_default, DefaultClause):
5✔
259
                self.add_literal_import("sqlalchemy", "text")
5✔
260

261
    def collect_imports_for_constraint(self, constraint: Constraint | Index) -> None:
5✔
262
        if isinstance(constraint, Index):
5✔
263
            if len(constraint.columns) > 1 or not uses_default_name(constraint):
5✔
264
                self.add_literal_import("sqlalchemy", "Index")
5✔
265
        elif isinstance(constraint, PrimaryKeyConstraint):
5✔
266
            if not uses_default_name(constraint):
5✔
267
                self.add_literal_import("sqlalchemy", "PrimaryKeyConstraint")
5✔
268
        elif isinstance(constraint, UniqueConstraint):
5✔
269
            if len(constraint.columns) > 1 or not uses_default_name(constraint):
5✔
270
                self.add_literal_import("sqlalchemy", "UniqueConstraint")
5✔
271
        elif isinstance(constraint, ForeignKeyConstraint):
5✔
272
            if len(constraint.columns) > 1 or not uses_default_name(constraint):
5✔
273
                self.add_literal_import("sqlalchemy", "ForeignKeyConstraint")
5✔
274
            else:
275
                self.add_import(ForeignKey)
5✔
276
        else:
277
            self.add_import(constraint)
5✔
278

279
    def add_import(self, obj: Any) -> None:
5✔
280
        # Don't store builtin imports
281
        if getattr(obj, "__module__", "builtins") == "builtins":
5✔
UNCOV
282
            return
×
283

284
        type_ = type(obj) if not isinstance(obj, type) else obj
5✔
285
        pkgname = type_.__module__
5✔
286

287
        # The column types have already been adapted towards generic types if possible,
288
        # so if this is still a vendor specific type (e.g., MySQL INTEGER) be sure to
289
        # use that rather than the generic sqlalchemy type as it might have different
290
        # constructor parameters.
291
        if pkgname.startswith("sqlalchemy.dialects."):
5✔
292
            dialect_pkgname = ".".join(pkgname.split(".")[0:3])
5✔
293
            dialect_pkg = import_module(dialect_pkgname)
5✔
294

295
            if type_.__name__ in dialect_pkg.__all__:
5✔
296
                pkgname = dialect_pkgname
5✔
297
        elif type_ is getattr(sqlalchemy, type_.__name__, None):
5✔
298
            pkgname = "sqlalchemy"
5✔
299
        else:
300
            pkgname = type_.__module__
5✔
301

302
        self.add_literal_import(pkgname, type_.__name__)
5✔
303

304
    def add_literal_import(self, pkgname: str, name: str) -> None:
5✔
305
        names = self.imports.setdefault(pkgname, set())
5✔
306
        names.add(name)
5✔
307

308
    def remove_literal_import(self, pkgname: str, name: str) -> None:
5✔
309
        names = self.imports.setdefault(pkgname, set())
5✔
310
        if name in names:
5✔
UNCOV
311
            names.remove(name)
×
312

313
    def add_module_import(self, pgkname: str) -> None:
5✔
314
        self.module_imports.add(pgkname)
5✔
315

316
    def group_imports(self) -> list[list[str]]:
5✔
317
        future_imports: list[str] = []
5✔
318
        stdlib_imports: list[str] = []
5✔
319
        thirdparty_imports: list[str] = []
5✔
320

321
        def get_collection(package: str) -> list[str]:
5✔
322
            collection = thirdparty_imports
5✔
323
            if package == "__future__":
5✔
UNCOV
324
                collection = future_imports
×
325
            elif package in self.stdlib_module_names:
5✔
326
                collection = stdlib_imports
5✔
327
            elif package in sys.modules:
5✔
328
                if "site-packages" not in (sys.modules[package].__file__ or ""):
5✔
329
                    collection = stdlib_imports
5✔
330
            return collection
5✔
331

332
        for package in sorted(self.imports):
5✔
333
            imports = ", ".join(sorted(self.imports[package]))
5✔
334

335
            collection = get_collection(package)
5✔
336
            collection.append(f"from {package} import {imports}")
5✔
337

338
        for module in sorted(self.module_imports):
5✔
339
            collection = get_collection(module)
5✔
340
            collection.append(f"import {module}")
5✔
341

342
        return [
5✔
343
            group
344
            for group in (future_imports, stdlib_imports, thirdparty_imports)
345
            if group
346
        ]
347

348
    def generate_models(self) -> list[Model]:
5✔
349
        models = [Model(table) for table in self.metadata.sorted_tables]
5✔
350

351
        # Collect the imports
352
        self.collect_imports(models)
5✔
353

354
        # Generate names for models
355
        global_names = {
5✔
356
            name for namespace in self.imports.values() for name in namespace
357
        }
358
        for model in models:
5✔
359
            self.generate_model_name(model, global_names)
5✔
360
            global_names.add(model.name)
5✔
361

362
        return models
5✔
363

364
    def generate_model_name(self, model: Model, global_names: set[str]) -> None:
5✔
365
        preferred_name = f"t_{model.table.name}"
5✔
366
        model.name = self.find_free_name(preferred_name, global_names)
5✔
367

368
    def render_module_variables(self, models: list[Model]) -> str:
5✔
369
        declarations = self.base.declarations
5✔
370

371
        if any(not isinstance(model, ModelClass) for model in models):
5✔
372
            if self.base.table_metadata_declaration is not None:
5✔
UNCOV
373
                declarations.append(self.base.table_metadata_declaration)
×
374

375
        return "\n".join(declarations)
5✔
376

377
    def render_models(self, models: list[Model]) -> str:
5✔
378
        rendered: list[str] = []
5✔
379
        for model in models:
5✔
380
            rendered_table = self.render_table(model.table)
5✔
381
            rendered.append(f"{model.name} = {rendered_table}")
5✔
382

383
        return "\n\n".join(rendered)
5✔
384

385
    def render_table(self, table: Table) -> str:
5✔
386
        args: list[str] = [f"{table.name!r}, {self.base.metadata_ref}"]
5✔
387
        kwargs: dict[str, object] = {}
5✔
388
        for column in table.columns:
5✔
389
            # Cast is required because of a bug in the SQLAlchemy stubs regarding
390
            # Table.columns
391
            args.append(self.render_column(column, True, is_table=True))
5✔
392

393
        for constraint in sorted(table.constraints, key=get_constraint_sort_key):
5✔
394
            if uses_default_name(constraint):
5✔
395
                if isinstance(constraint, PrimaryKeyConstraint):
5✔
396
                    continue
5✔
397
                elif isinstance(constraint, (ForeignKeyConstraint, UniqueConstraint)):
5✔
398
                    if len(constraint.columns) == 1:
5✔
399
                        continue
5✔
400

401
            args.append(self.render_constraint(constraint))
5✔
402

403
        for index in sorted(table.indexes, key=lambda i: cast(str, i.name)):
5✔
404
            # One-column indexes should be rendered as index=True on columns
405
            if len(index.columns) > 1 or not uses_default_name(index):
5✔
406
                args.append(self.render_index(index))
5✔
407

408
        if table.schema:
5✔
409
            kwargs["schema"] = repr(table.schema)
5✔
410

411
        table_comment = getattr(table, "comment", None)
5✔
412
        if table_comment:
5✔
413
            kwargs["comment"] = repr(table.comment)
5✔
414

415
        # add info + dialect kwargs for callable context (opt-in)
416
        if self.include_dialect_options_and_info:
5✔
417
            self._add_dialect_kwargs_and_info(table, kwargs, values_for_dict=False)
5✔
418

419
        return render_callable("Table", *args, kwargs=kwargs, indentation="    ")
5✔
420

421
    def render_index(self, index: Index) -> str:
5✔
422
        extra_args = [repr(col.name) for col in index.columns]
5✔
423
        kwargs = {}
5✔
424
        if index.unique:
5✔
425
            kwargs["unique"] = True
5✔
426

427
        return render_callable("Index", repr(index.name), *extra_args, kwargs=kwargs)
5✔
428

429
    # TODO find better solution for is_table
430
    def render_column(
5✔
431
        self, column: Column[Any], show_name: bool, is_table: bool = False
432
    ) -> str:
433
        args = []
5✔
434
        kwargs: dict[str, Any] = {}
5✔
435
        kwarg = []
5✔
436
        is_part_of_composite_pk = (
5✔
437
            column.primary_key and len(column.table.primary_key) > 1
438
        )
439
        dedicated_fks = [
5✔
440
            c
441
            for c in column.foreign_keys
442
            if c.constraint
443
            and len(c.constraint.columns) == 1
444
            and uses_default_name(c.constraint)
445
        ]
446
        is_unique = any(
5✔
447
            isinstance(c, UniqueConstraint)
448
            and set(c.columns) == {column}
449
            and uses_default_name(c)
450
            for c in column.table.constraints
451
        )
452
        is_unique = is_unique or any(
5✔
453
            i.unique and set(i.columns) == {column} and uses_default_name(i)
454
            for i in column.table.indexes
455
        )
456
        is_primary = (
5✔
457
            any(
458
                isinstance(c, PrimaryKeyConstraint)
459
                and column.name in c.columns
460
                and uses_default_name(c)
461
                for c in column.table.constraints
462
            )
463
            or column.primary_key
464
        )
465
        has_index = any(
5✔
466
            set(i.columns) == {column} and uses_default_name(i)
467
            for i in column.table.indexes
468
        )
469

470
        if show_name:
5✔
471
            args.append(repr(column.name))
5✔
472

473
        # Render the column type if there are no foreign keys on it or any of them
474
        # points back to itself
475
        if not dedicated_fks or any(fk.column is column for fk in dedicated_fks):
5✔
476
            args.append(self.render_column_type(column.type))
5✔
477

478
        for fk in dedicated_fks:
5✔
479
            args.append(self.render_constraint(fk))
5✔
480

481
        if column.default:
5✔
482
            args.append(repr(column.default))
5✔
483

484
        if column.key != column.name:
5✔
UNCOV
485
            kwargs["key"] = column.key
×
486
        if is_primary:
5✔
487
            kwargs["primary_key"] = True
5✔
488
        if not column.nullable and not column.primary_key:
5✔
489
            kwargs["nullable"] = False
5✔
490
        if column.nullable and is_part_of_composite_pk:
5✔
491
            kwargs["nullable"] = True
5✔
492

493
        if is_unique:
5✔
494
            column.unique = True
5✔
495
            kwargs["unique"] = True
5✔
496
        if has_index:
5✔
497
            column.index = True
5✔
498
            kwarg.append("index")
5✔
499
            kwargs["index"] = True
5✔
500

501
        if isinstance(column.server_default, DefaultClause):
5✔
502
            kwargs["server_default"] = render_callable(
5✔
503
                "text", repr(cast(TextClause, column.server_default.arg).text)
504
            )
505
        elif isinstance(column.server_default, Computed):
5✔
506
            expression = str(column.server_default.sqltext)
5✔
507

508
            computed_kwargs = {}
5✔
509
            if column.server_default.persisted is not None:
5✔
510
                computed_kwargs["persisted"] = column.server_default.persisted
5✔
511

512
            args.append(
5✔
513
                render_callable("Computed", repr(expression), kwargs=computed_kwargs)
514
            )
515
        elif isinstance(column.server_default, Identity):
5✔
516
            args.append(repr(column.server_default))
5✔
517
        elif column.server_default:
5✔
518
            kwargs["server_default"] = repr(column.server_default)
×
519

520
        comment = getattr(column, "comment", None)
5✔
521
        if comment:
5✔
522
            kwargs["comment"] = repr(comment)
5✔
523

524
        # add column info + dialect kwargs for callable context (opt-in)
525
        if self.include_dialect_options_and_info:
5✔
526
            self._add_dialect_kwargs_and_info(column, kwargs, values_for_dict=False)
5✔
527

528
        return self.render_column_callable(is_table, *args, **kwargs)
5✔
529

530
    def render_column_callable(self, is_table: bool, *args: Any, **kwargs: Any) -> str:
5✔
531
        if is_table:
5✔
532
            self.add_import(Column)
5✔
533
            return render_callable("Column", *args, kwargs=kwargs)
5✔
534
        else:
535
            return render_callable("mapped_column", *args, kwargs=kwargs)
5✔
536

537
    def render_column_type(self, coltype: TypeEngine[Any]) -> str:
5✔
538
        args = []
5✔
539
        kwargs: dict[str, Any] = {}
5✔
540
        sig = inspect.signature(coltype.__class__.__init__)
5✔
541
        defaults = {param.name: param.default for param in sig.parameters.values()}
5✔
542
        missing = object()
5✔
543
        use_kwargs = False
5✔
544
        for param in list(sig.parameters.values())[1:]:
5✔
545
            # Remove annoyances like _warn_on_bytestring
546
            if param.name.startswith("_"):
5✔
547
                continue
5✔
548
            elif param.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
5✔
549
                use_kwargs = True
5✔
550
                continue
5✔
551

552
            value = getattr(coltype, param.name, missing)
5✔
553

554
            if isinstance(value, (JSONB, JSON)):
5✔
555
                # Remove astext_type if it's the default
556
                if (
5✔
557
                    isinstance(value.astext_type, Text)
558
                    and value.astext_type.length is None
559
                ):
560
                    value.astext_type = None  # type: ignore[assignment]
5✔
561
                else:
562
                    self.add_import(Text)
5✔
563

564
            default = defaults.get(param.name, missing)
5✔
565
            if isinstance(value, TextClause):
5✔
566
                self.add_literal_import("sqlalchemy", "text")
5✔
567
                rendered_value = render_callable("text", repr(value.text))
5✔
568
            else:
569
                rendered_value = repr(value)
5✔
570

571
            if value is missing or value == default:
5✔
572
                use_kwargs = True
5✔
573
            elif use_kwargs:
5✔
574
                kwargs[param.name] = rendered_value
5✔
575
            else:
576
                args.append(rendered_value)
5✔
577

578
        vararg = next(
5✔
579
            (
580
                param.name
581
                for param in sig.parameters.values()
582
                if param.kind is Parameter.VAR_POSITIONAL
583
            ),
584
            None,
585
        )
586
        if vararg and hasattr(coltype, vararg):
5✔
587
            varargs_repr = [repr(arg) for arg in getattr(coltype, vararg)]
5✔
588
            args.extend(varargs_repr)
5✔
589

590
        # These arguments cannot be autodetected from the Enum initializer
591
        if isinstance(coltype, Enum):
5✔
592
            for colname in "name", "schema":
5✔
593
                if (value := getattr(coltype, colname)) is not None:
5✔
594
                    kwargs[colname] = repr(value)
5✔
595

596
        if isinstance(coltype, (JSONB, JSON)):
5✔
597
            # Remove astext_type if it's the default
598
            if (
5✔
599
                isinstance(coltype.astext_type, Text)
600
                and coltype.astext_type.length is None
601
            ):
602
                del kwargs["astext_type"]
5✔
603

604
        if args or kwargs:
5✔
605
            return render_callable(coltype.__class__.__name__, *args, kwargs=kwargs)
5✔
606
        else:
607
            return coltype.__class__.__name__
5✔
608

609
    def render_constraint(self, constraint: Constraint | ForeignKey) -> str:
5✔
610
        def add_fk_options(*opts: Any) -> None:
5✔
611
            args.extend(repr(opt) for opt in opts)
5✔
612
            for attr in "ondelete", "onupdate", "deferrable", "initially", "match":
5✔
613
                value = getattr(constraint, attr, None)
5✔
614
                if value:
5✔
615
                    kwargs[attr] = repr(value)
5✔
616

617
        args: list[str] = []
5✔
618
        kwargs: dict[str, Any] = {}
5✔
619
        if isinstance(constraint, ForeignKey):
5✔
620
            remote_column = (
5✔
621
                f"{constraint.column.table.fullname}.{constraint.column.name}"
622
            )
623
            add_fk_options(remote_column)
5✔
624
        elif isinstance(constraint, ForeignKeyConstraint):
5✔
625
            local_columns = get_column_names(constraint)
5✔
626
            remote_columns = [
5✔
627
                f"{fk.column.table.fullname}.{fk.column.name}"
628
                for fk in constraint.elements
629
            ]
630
            add_fk_options(local_columns, remote_columns)
5✔
631
        elif isinstance(constraint, CheckConstraint):
5✔
632
            args.append(repr(get_compiled_expression(constraint.sqltext, self.bind)))
5✔
633
        elif isinstance(constraint, (UniqueConstraint, PrimaryKeyConstraint)):
5✔
634
            args.extend(repr(col.name) for col in constraint.columns)
5✔
635
        else:
UNCOV
636
            raise TypeError(
×
637
                f"Cannot render constraint of type {constraint.__class__.__name__}"
638
            )
639

640
        if isinstance(constraint, Constraint) and not uses_default_name(constraint):
5✔
641
            kwargs["name"] = repr(constraint.name)
5✔
642

643
        return render_callable(constraint.__class__.__name__, *args, kwargs=kwargs)
5✔
644

645
    def _add_dialect_kwargs_and_info(
5✔
646
        self, obj: Any, target_kwargs: dict[str, object], *, values_for_dict: bool
647
    ) -> None:
648
        """
649
        Merge SchemaItem-like object's .info and .dialect_kwargs into target_kwargs.
650
        - values_for_dict=True: keep raw values so pretty-printer emits repr() (for __table_args__ dict)
651
        - values_for_dict=False: set values to repr() strings (for callable kwargs)
652
        """
653
        info_dict = getattr(obj, "info", None)
5✔
654
        if info_dict:
5✔
655
            target_kwargs["info"] = info_dict if values_for_dict else repr(info_dict)
5✔
656

657
        dialect_keys: list[str]
658
        try:
5✔
659
            dialect_keys = sorted(getattr(obj, "dialect_kwargs"))
5✔
NEW
660
        except Exception:
×
NEW
661
            return
×
662

663
        dialect_kwargs = getattr(obj, "dialect_kwargs", {})
5✔
664
        for key in dialect_keys:
5✔
665
            try:
5✔
666
                value = dialect_kwargs[key]
5✔
NEW
667
            except Exception:
×
NEW
668
                continue
×
669
            # Render values:
670
            # - callable context (values_for_dict=False): produce a string expression.
671
            #   primitives use repr(value); custom objects stringify then repr().
672
            # - dict context (values_for_dict=True): pass raw primitives / str;
673
            #   custom objects become str(value) so pformat quotes them.
674
            if values_for_dict:
5✔
675
                if isinstance(value, type(None) | bool | int | float):
5✔
NEW
676
                    target_kwargs[key] = value
×
677
                elif isinstance(value, str | dict | list):
5✔
678
                    target_kwargs[key] = value
5✔
679
                else:
680
                    target_kwargs[key] = str(value)
5✔
681
            else:
682
                if isinstance(
5✔
683
                    value, type(None) | bool | int | float | str | dict | list
684
                ):
685
                    target_kwargs[key] = repr(value)
5✔
686
                else:
687
                    target_kwargs[key] = repr(str(value))
5✔
688

689
    def should_ignore_table(self, table: Table) -> bool:
5✔
690
        # Support for Alembic and sqlalchemy-migrate -- never expose the schema version
691
        # tables
692
        return table.name in ("alembic_version", "migrate_version")
5✔
693

694
    def find_free_name(
5✔
695
        self, name: str, global_names: set[str], local_names: Collection[str] = ()
696
    ) -> str:
697
        """
698
        Generate an attribute name that does not clash with other local or global names.
699
        """
700
        name = name.strip()
5✔
701
        assert name, "Identifier cannot be empty"
5✔
702
        name = _re_invalid_identifier.sub("_", name)
5✔
703
        if name[0].isdigit():
5✔
704
            name = "_" + name
5✔
705
        elif iskeyword(name) or name == "metadata":
5✔
706
            name += "_"
5✔
707

708
        original = name
5✔
709
        for i in count():
5✔
710
            if name not in global_names and name not in local_names:
5✔
711
                break
5✔
712

713
            name = original + (str(i) if i else "_")
5✔
714

715
        return name
5✔
716

717
    def fix_column_types(self, table: Table) -> None:
5✔
718
        """Adjust the reflected column types."""
719
        # Detect check constraints for boolean and enum columns
720
        for constraint in table.constraints.copy():
5✔
721
            if isinstance(constraint, CheckConstraint):
5✔
722
                sqltext = get_compiled_expression(constraint.sqltext, self.bind)
5✔
723

724
                # Turn any integer-like column with a CheckConstraint like
725
                # "column IN (0, 1)" into a Boolean
726
                match = _re_boolean_check_constraint.match(sqltext)
5✔
727
                if match:
5✔
728
                    colname_match = _re_column_name.match(match.group(1))
5✔
729
                    if colname_match:
5✔
730
                        colname = colname_match.group(3)
5✔
731
                        table.constraints.remove(constraint)
5✔
732
                        table.c[colname].type = Boolean()
5✔
733
                        continue
5✔
734

735
                # Turn any string-type column with a CheckConstraint like
736
                # "column IN (...)" into an Enum
737
                match = _re_enum_check_constraint.match(sqltext)
5✔
738
                if match:
5✔
739
                    colname_match = _re_column_name.match(match.group(1))
5✔
740
                    if colname_match:
5✔
741
                        colname = colname_match.group(3)
5✔
742
                        items = match.group(2)
5✔
743
                        if isinstance(table.c[colname].type, String):
5✔
744
                            table.constraints.remove(constraint)
5✔
745
                            if not isinstance(table.c[colname].type, Enum):
5✔
746
                                options = _re_enum_item.findall(items)
5✔
747
                                table.c[colname].type = Enum(
5✔
748
                                    *options, native_enum=False
749
                                )
750

751
                            continue
5✔
752

753
        for column in table.c:
5✔
754
            if not self.keep_dialect_types:
5✔
755
                try:
5✔
756
                    column.type = self.get_adapted_type(column.type)
5✔
757
                except CompileError:
5✔
758
                    continue
5✔
759

760
            # PostgreSQL specific fix: detect sequences from server_default
761
            if column.server_default and self.bind.dialect.name == "postgresql":
5✔
762
                if isinstance(column.server_default, DefaultClause) and isinstance(
5✔
763
                    column.server_default.arg, TextClause
764
                ):
765
                    schema, seqname = decode_postgresql_sequence(
5✔
766
                        column.server_default.arg
767
                    )
768
                    if seqname:
5✔
769
                        # Add an explicit sequence
770
                        if seqname != f"{column.table.name}_{column.name}_seq":
5✔
771
                            column.default = sqlalchemy.Sequence(seqname, schema=schema)
5✔
772

773
                        column.server_default = None
5✔
774

775
    def get_adapted_type(self, coltype: Any) -> Any:
5✔
776
        compiled_type = coltype.compile(self.bind.engine.dialect)
5✔
777
        for supercls in coltype.__class__.__mro__:
5✔
778
            if not supercls.__name__.startswith("_") and hasattr(
5✔
779
                supercls, "__visit_name__"
780
            ):
781
                # Don't try to adapt UserDefinedType as it's not a proper column type
782
                if supercls is UserDefinedType or issubclass(supercls, TypeDecorator):
5✔
783
                    return coltype
5✔
784

785
                # Hack to fix adaptation of the Enum class which is broken since
786
                # SQLAlchemy 1.2
787
                kw = {}
5✔
788
                if supercls is Enum:
5✔
789
                    kw["name"] = coltype.name
5✔
790
                    if coltype.schema:
5✔
791
                        kw["schema"] = coltype.schema
5✔
792

793
                # Hack to fix Postgres DOMAIN type adaptation, broken as of SQLAlchemy 2.0.42
794
                # For additional information - https://github.com/agronholm/sqlacodegen/issues/416#issuecomment-3417480599
795
                if supercls is DOMAIN:
5✔
796
                    if coltype.default:
5✔
UNCOV
797
                        kw["default"] = coltype.default
×
798
                    if coltype.constraint_name is not None:
5✔
799
                        kw["constraint_name"] = coltype.constraint_name
5✔
800
                    if coltype.not_null:
5✔
UNCOV
801
                        kw["not_null"] = coltype.not_null
×
802
                    if coltype.check is not None:
5✔
803
                        kw["check"] = coltype.check
5✔
804
                    if coltype.create_type:
5✔
805
                        kw["create_type"] = coltype.create_type
5✔
806

807
                try:
5✔
808
                    new_coltype = coltype.adapt(supercls)
5✔
809
                except TypeError:
5✔
810
                    # If the adaptation fails, don't try again
811
                    break
5✔
812

813
                for key, value in kw.items():
5✔
814
                    setattr(new_coltype, key, value)
5✔
815

816
                if isinstance(coltype, ARRAY):
5✔
817
                    new_coltype.item_type = self.get_adapted_type(new_coltype.item_type)
5✔
818

819
                try:
5✔
820
                    # If the adapted column type does not render the same as the
821
                    # original, don't substitute it
822
                    if new_coltype.compile(self.bind.engine.dialect) != compiled_type:
5✔
823
                        break
5✔
824
                except CompileError:
5✔
825
                    # If the adapted column type can't be compiled, don't substitute it
826
                    break
5✔
827

828
                # Stop on the first valid non-uppercase column type class
829
                coltype = new_coltype
5✔
830
                if supercls.__name__ != supercls.__name__.upper():
5✔
831
                    break
5✔
832

833
        return coltype
5✔
834

835

836
class DeclarativeGenerator(TablesGenerator):
5✔
837
    valid_options: ClassVar[set[str]] = TablesGenerator.valid_options | {
5✔
838
        "use_inflect",
839
        "nojoined",
840
        "nobidi",
841
        "noidsuffix",
842
    }
843

844
    def __init__(
5✔
845
        self,
846
        metadata: MetaData,
847
        bind: Connection | Engine,
848
        options: Sequence[str],
849
        *,
850
        indentation: str = "    ",
851
        base_class_name: str = "Base",
852
    ):
853
        super().__init__(metadata, bind, options, indentation=indentation)
5✔
854
        self.base_class_name: str = base_class_name
5✔
855
        self.inflect_engine = inflect.engine()
5✔
856

857
    def generate_base(self) -> None:
5✔
858
        self.base = Base(
5✔
859
            literal_imports=[LiteralImport("sqlalchemy.orm", "DeclarativeBase")],
860
            declarations=[
861
                f"class {self.base_class_name}(DeclarativeBase):",
862
                f"{self.indentation}pass",
863
            ],
864
            metadata_ref=f"{self.base_class_name}.metadata",
865
        )
866

867
    def collect_imports(self, models: Iterable[Model]) -> None:
5✔
868
        super().collect_imports(models)
5✔
869
        if any(isinstance(model, ModelClass) for model in models):
5✔
870
            self.add_literal_import("sqlalchemy.orm", "Mapped")
5✔
871
            self.add_literal_import("sqlalchemy.orm", "mapped_column")
5✔
872

873
    def collect_imports_for_model(self, model: Model) -> None:
5✔
874
        super().collect_imports_for_model(model)
5✔
875
        if isinstance(model, ModelClass):
5✔
876
            if model.relationships:
5✔
877
                self.add_literal_import("sqlalchemy.orm", "relationship")
5✔
878

879
    def generate_models(self) -> list[Model]:
5✔
880
        models_by_table_name: dict[str, Model] = {}
5✔
881

882
        # Pick association tables from the metadata into their own set, don't process
883
        # them normally
884
        links: defaultdict[str, list[Model]] = defaultdict(lambda: [])
5✔
885
        for table in self.metadata.sorted_tables:
5✔
886
            qualified_name = qualified_table_name(table)
5✔
887

888
            # Link tables have exactly two foreign key constraints and all columns are
889
            # involved in them
890
            fk_constraints = sorted(
5✔
891
                table.foreign_key_constraints, key=get_constraint_sort_key
892
            )
893
            if len(fk_constraints) == 2 and all(
5✔
894
                col.foreign_keys for col in table.columns
895
            ):
896
                model = models_by_table_name[qualified_name] = Model(table)
5✔
897
                tablename = fk_constraints[0].elements[0].column.table.name
5✔
898
                links[tablename].append(model)
5✔
899
                continue
5✔
900

901
            # Only form model classes for tables that have a primary key and are not
902
            # association tables
903
            if not table.primary_key:
5✔
904
                models_by_table_name[qualified_name] = Model(table)
5✔
905
            else:
906
                model = ModelClass(table)
5✔
907
                models_by_table_name[qualified_name] = model
5✔
908

909
                # Fill in the columns
910
                for column in table.c:
5✔
911
                    column_attr = ColumnAttribute(model, column)
5✔
912
                    model.columns.append(column_attr)
5✔
913

914
        # Add relationships
915
        for model in models_by_table_name.values():
5✔
916
            if isinstance(model, ModelClass):
5✔
917
                self.generate_relationships(
5✔
918
                    model, models_by_table_name, links[model.table.name]
919
                )
920

921
        # Nest inherited classes in their superclasses to ensure proper ordering
922
        if "nojoined" not in self.options:
5✔
923
            for model in list(models_by_table_name.values()):
5✔
924
                if not isinstance(model, ModelClass):
5✔
925
                    continue
5✔
926

927
                pk_column_names = {col.name for col in model.table.primary_key.columns}
5✔
928
                for constraint in model.table.foreign_key_constraints:
5✔
929
                    if set(get_column_names(constraint)) == pk_column_names:
5✔
930
                        target = models_by_table_name[
5✔
931
                            qualified_table_name(constraint.elements[0].column.table)
932
                        ]
933
                        if isinstance(target, ModelClass):
5✔
934
                            model.parent_class = target
5✔
935
                            target.children.append(model)
5✔
936

937
        # Change base if we only have tables
938
        if not any(
5✔
939
            isinstance(model, ModelClass) for model in models_by_table_name.values()
940
        ):
941
            super().generate_base()
5✔
942

943
        # Collect the imports
944
        self.collect_imports(models_by_table_name.values())
5✔
945

946
        # Rename models and their attributes that conflict with imports or other
947
        # attributes
948
        global_names = {
5✔
949
            name for namespace in self.imports.values() for name in namespace
950
        }
951
        for model in models_by_table_name.values():
5✔
952
            self.generate_model_name(model, global_names)
5✔
953
            global_names.add(model.name)
5✔
954

955
        return list(models_by_table_name.values())
5✔
956

957
    def generate_relationships(
5✔
958
        self,
959
        source: ModelClass,
960
        models_by_table_name: dict[str, Model],
961
        association_tables: list[Model],
962
    ) -> list[RelationshipAttribute]:
963
        relationships: list[RelationshipAttribute] = []
5✔
964
        reverse_relationship: RelationshipAttribute | None
965

966
        # Add many-to-one (and one-to-many) relationships
967
        pk_column_names = {col.name for col in source.table.primary_key.columns}
5✔
968
        for constraint in sorted(
5✔
969
            source.table.foreign_key_constraints, key=get_constraint_sort_key
970
        ):
971
            target = models_by_table_name[
5✔
972
                qualified_table_name(constraint.elements[0].column.table)
973
            ]
974
            if isinstance(target, ModelClass):
5✔
975
                if "nojoined" not in self.options:
5✔
976
                    if set(get_column_names(constraint)) == pk_column_names:
5✔
977
                        parent = models_by_table_name[
5✔
978
                            qualified_table_name(constraint.elements[0].column.table)
979
                        ]
980
                        if isinstance(parent, ModelClass):
5✔
981
                            source.parent_class = parent
5✔
982
                            parent.children.append(source)
5✔
983
                            continue
5✔
984

985
                # Add uselist=False to One-to-One relationships
986
                column_names = get_column_names(constraint)
5✔
987
                if any(
5✔
988
                    isinstance(c, (PrimaryKeyConstraint, UniqueConstraint))
989
                    and {col.name for col in c.columns} == set(column_names)
990
                    for c in constraint.table.constraints
991
                ):
992
                    r_type = RelationshipType.ONE_TO_ONE
5✔
993
                else:
994
                    r_type = RelationshipType.MANY_TO_ONE
5✔
995

996
                relationship = RelationshipAttribute(r_type, source, target, constraint)
5✔
997
                source.relationships.append(relationship)
5✔
998

999
                # For self referential relationships, remote_side needs to be set
1000
                if source is target:
5✔
1001
                    relationship.remote_side = [
5✔
1002
                        source.get_column_attribute(col.name)
1003
                        for col in constraint.referred_table.primary_key
1004
                    ]
1005

1006
                # If the two tables share more than one foreign key constraint,
1007
                # SQLAlchemy needs an explicit primaryjoin to figure out which column(s)
1008
                # it needs
1009
                common_fk_constraints = get_common_fk_constraints(
5✔
1010
                    source.table, target.table
1011
                )
1012
                if len(common_fk_constraints) > 1:
5✔
1013
                    relationship.foreign_keys = [
5✔
1014
                        source.get_column_attribute(key)
1015
                        for key in constraint.column_keys
1016
                    ]
1017

1018
                # Generate the opposite end of the relationship in the target class
1019
                if "nobidi" not in self.options:
5✔
1020
                    if r_type is RelationshipType.MANY_TO_ONE:
5✔
1021
                        r_type = RelationshipType.ONE_TO_MANY
5✔
1022

1023
                    reverse_relationship = RelationshipAttribute(
5✔
1024
                        r_type,
1025
                        target,
1026
                        source,
1027
                        constraint,
1028
                        foreign_keys=relationship.foreign_keys,
1029
                        backref=relationship,
1030
                    )
1031
                    relationship.backref = reverse_relationship
5✔
1032
                    target.relationships.append(reverse_relationship)
5✔
1033

1034
                    # For self referential relationships, remote_side needs to be set
1035
                    if source is target:
5✔
1036
                        reverse_relationship.remote_side = [
5✔
1037
                            source.get_column_attribute(colname)
1038
                            for colname in constraint.column_keys
1039
                        ]
1040

1041
        # Add many-to-many relationships
1042
        for association_table in association_tables:
5✔
1043
            fk_constraints = sorted(
5✔
1044
                association_table.table.foreign_key_constraints,
1045
                key=get_constraint_sort_key,
1046
            )
1047
            target = models_by_table_name[
5✔
1048
                qualified_table_name(fk_constraints[1].elements[0].column.table)
1049
            ]
1050
            if isinstance(target, ModelClass):
5✔
1051
                relationship = RelationshipAttribute(
5✔
1052
                    RelationshipType.MANY_TO_MANY,
1053
                    source,
1054
                    target,
1055
                    fk_constraints[1],
1056
                    association_table,
1057
                )
1058
                source.relationships.append(relationship)
5✔
1059

1060
                # Generate the opposite end of the relationship in the target class
1061
                reverse_relationship = None
5✔
1062
                if "nobidi" not in self.options:
5✔
1063
                    reverse_relationship = RelationshipAttribute(
5✔
1064
                        RelationshipType.MANY_TO_MANY,
1065
                        target,
1066
                        source,
1067
                        fk_constraints[0],
1068
                        association_table,
1069
                        relationship,
1070
                    )
1071
                    relationship.backref = reverse_relationship
5✔
1072
                    target.relationships.append(reverse_relationship)
5✔
1073

1074
                # Add a primary/secondary join for self-referential many-to-many
1075
                # relationships
1076
                if source is target:
5✔
1077
                    both_relationships = [relationship]
5✔
1078
                    reverse_flags = [False, True]
5✔
1079
                    if reverse_relationship:
5✔
1080
                        both_relationships.append(reverse_relationship)
5✔
1081

1082
                    for relationship, reverse in zip(both_relationships, reverse_flags):
5✔
1083
                        if (
5✔
1084
                            not relationship.association_table
1085
                            or not relationship.constraint
1086
                        ):
UNCOV
1087
                            continue
×
1088

1089
                        constraints = sorted(
5✔
1090
                            relationship.constraint.table.foreign_key_constraints,
1091
                            key=get_constraint_sort_key,
1092
                            reverse=reverse,
1093
                        )
1094
                        pri_pairs = zip(
5✔
1095
                            get_column_names(constraints[0]), constraints[0].elements
1096
                        )
1097
                        sec_pairs = zip(
5✔
1098
                            get_column_names(constraints[1]), constraints[1].elements
1099
                        )
1100
                        relationship.primaryjoin = [
5✔
1101
                            (
1102
                                relationship.source,
1103
                                elem.column.name,
1104
                                relationship.association_table,
1105
                                col,
1106
                            )
1107
                            for col, elem in pri_pairs
1108
                        ]
1109
                        relationship.secondaryjoin = [
5✔
1110
                            (
1111
                                relationship.target,
1112
                                elem.column.name,
1113
                                relationship.association_table,
1114
                                col,
1115
                            )
1116
                            for col, elem in sec_pairs
1117
                        ]
1118

1119
        return relationships
5✔
1120

1121
    def generate_model_name(self, model: Model, global_names: set[str]) -> None:
5✔
1122
        if isinstance(model, ModelClass):
5✔
1123
            preferred_name = _re_invalid_identifier.sub("_", model.table.name)
5✔
1124
            preferred_name = "".join(
5✔
1125
                part[:1].upper() + part[1:] for part in preferred_name.split("_")
1126
            )
1127
            if "use_inflect" in self.options:
5✔
1128
                singular_name = self.inflect_engine.singular_noun(preferred_name)
5✔
1129
                if singular_name:
5✔
1130
                    preferred_name = singular_name
5✔
1131

1132
            model.name = self.find_free_name(preferred_name, global_names)
5✔
1133

1134
            # Fill in the names for column attributes
1135
            local_names: set[str] = set()
5✔
1136
            for column_attr in model.columns:
5✔
1137
                self.generate_column_attr_name(column_attr, global_names, local_names)
5✔
1138
                local_names.add(column_attr.name)
5✔
1139

1140
            # Fill in the names for relationship attributes
1141
            for relationship in model.relationships:
5✔
1142
                self.generate_relationship_name(relationship, global_names, local_names)
5✔
1143
                local_names.add(relationship.name)
5✔
1144
        else:
1145
            super().generate_model_name(model, global_names)
5✔
1146

1147
    def generate_column_attr_name(
5✔
1148
        self,
1149
        column_attr: ColumnAttribute,
1150
        global_names: set[str],
1151
        local_names: set[str],
1152
    ) -> None:
1153
        column_attr.name = self.find_free_name(
5✔
1154
            column_attr.column.name, global_names, local_names
1155
        )
1156

1157
    def generate_relationship_name(
5✔
1158
        self,
1159
        relationship: RelationshipAttribute,
1160
        global_names: set[str],
1161
        local_names: set[str],
1162
    ) -> None:
1163
        # Self referential reverse relationships
1164
        preferred_name: str
1165
        if (
5✔
1166
            relationship.type
1167
            in (RelationshipType.ONE_TO_MANY, RelationshipType.ONE_TO_ONE)
1168
            and relationship.source is relationship.target
1169
            and relationship.backref
1170
            and relationship.backref.name
1171
        ):
1172
            preferred_name = relationship.backref.name + "_reverse"
5✔
1173
        else:
1174
            preferred_name = relationship.target.table.name
5✔
1175

1176
            # If there's a constraint with a single column that ends with "_id", use the
1177
            # preceding part as the relationship name
1178
            if relationship.constraint and "noidsuffix" not in self.options:
5✔
1179
                is_source = relationship.source.table is relationship.constraint.table
5✔
1180
                if is_source or relationship.type not in (
5✔
1181
                    RelationshipType.ONE_TO_ONE,
1182
                    RelationshipType.ONE_TO_MANY,
1183
                ):
1184
                    column_names = [c.name for c in relationship.constraint.columns]
5✔
1185
                    if len(column_names) == 1 and column_names[0].endswith("_id"):
5✔
1186
                        preferred_name = column_names[0][:-3]
5✔
1187

1188
            if "use_inflect" in self.options:
5✔
1189
                inflected_name: str | Literal[False]
1190
                if relationship.type in (
5✔
1191
                    RelationshipType.ONE_TO_MANY,
1192
                    RelationshipType.MANY_TO_MANY,
1193
                ):
1194
                    if not self.inflect_engine.singular_noun(preferred_name):
5✔
UNCOV
1195
                        preferred_name = self.inflect_engine.plural_noun(preferred_name)
×
1196
                else:
1197
                    inflected_name = self.inflect_engine.singular_noun(preferred_name)
5✔
1198
                    if inflected_name:
5✔
1199
                        preferred_name = inflected_name
5✔
1200

1201
        relationship.name = self.find_free_name(
5✔
1202
            preferred_name, global_names, local_names
1203
        )
1204

1205
    def render_models(self, models: list[Model]) -> str:
5✔
1206
        rendered: list[str] = []
5✔
1207
        for model in models:
5✔
1208
            if isinstance(model, ModelClass):
5✔
1209
                rendered.append(self.render_class(model))
5✔
1210
            else:
1211
                rendered.append(f"{model.name} = {self.render_table(model.table)}")
5✔
1212

1213
        return "\n\n\n".join(rendered)
5✔
1214

1215
    def render_class(self, model: ModelClass) -> str:
5✔
1216
        sections: list[str] = []
5✔
1217

1218
        # Render class variables / special declarations
1219
        class_vars: str = self.render_class_variables(model)
5✔
1220
        if class_vars:
5✔
1221
            sections.append(class_vars)
5✔
1222

1223
        # Render column attributes
1224
        rendered_column_attributes: list[str] = []
5✔
1225
        for nullable in (False, True):
5✔
1226
            for column_attr in model.columns:
5✔
1227
                if column_attr.column.nullable is nullable:
5✔
1228
                    rendered_column_attributes.append(
5✔
1229
                        self.render_column_attribute(column_attr)
1230
                    )
1231

1232
        if rendered_column_attributes:
5✔
1233
            sections.append("\n".join(rendered_column_attributes))
5✔
1234

1235
        # Render relationship attributes
1236
        rendered_relationship_attributes: list[str] = [
5✔
1237
            self.render_relationship(relationship)
1238
            for relationship in model.relationships
1239
        ]
1240

1241
        if rendered_relationship_attributes:
5✔
1242
            sections.append("\n".join(rendered_relationship_attributes))
5✔
1243

1244
        declaration = self.render_class_declaration(model)
5✔
1245
        rendered_sections = "\n\n".join(
5✔
1246
            indent(section, self.indentation) for section in sections
1247
        )
1248
        return f"{declaration}\n{rendered_sections}"
5✔
1249

1250
    def render_class_declaration(self, model: ModelClass) -> str:
5✔
1251
        parent_class_name = (
5✔
1252
            model.parent_class.name if model.parent_class else self.base_class_name
1253
        )
1254
        return f"class {model.name}({parent_class_name}):"
5✔
1255

1256
    def render_class_variables(self, model: ModelClass) -> str:
5✔
1257
        variables = [f"__tablename__ = {model.table.name!r}"]
5✔
1258

1259
        # Render constraints and indexes as __table_args__
1260
        table_args = self.render_table_args(model.table)
5✔
1261
        if table_args:
5✔
1262
            variables.append(f"__table_args__ = {table_args}")
5✔
1263

1264
        return "\n".join(variables)
5✔
1265

1266
    def render_table_args(self, table: Table) -> str:
5✔
1267
        args: list[str] = []
5✔
1268
        kwargs: dict[str, object] = {}
5✔
1269

1270
        # Render constraints
1271
        for constraint in sorted(table.constraints, key=get_constraint_sort_key):
5✔
1272
            if uses_default_name(constraint):
5✔
1273
                if isinstance(constraint, PrimaryKeyConstraint):
5✔
1274
                    continue
5✔
1275
                if (
5✔
1276
                    isinstance(constraint, (ForeignKeyConstraint, UniqueConstraint))
1277
                    and len(constraint.columns) == 1
1278
                ):
1279
                    continue
5✔
1280

1281
            args.append(self.render_constraint(constraint))
5✔
1282

1283
        # Render indexes
1284
        for index in sorted(table.indexes, key=lambda i: cast(str, i.name)):
5✔
1285
            if len(index.columns) > 1 or not uses_default_name(index):
5✔
1286
                args.append(self.render_index(index))
5✔
1287

1288
        if table.schema:
5✔
1289
            kwargs["schema"] = table.schema
5✔
1290

1291
        if table.comment:
5✔
1292
            kwargs["comment"] = table.comment
5✔
1293

1294
        # add info + dialect kwargs for dict context (__table_args__) (opt-in)
1295
        if self.include_dialect_options_and_info:
5✔
1296
            self._add_dialect_kwargs_and_info(table, kwargs, values_for_dict=True)
5✔
1297

1298
        if kwargs:
5✔
1299
            formatted_kwargs = pformat(kwargs)
5✔
1300
            if not args:
5✔
1301
                return formatted_kwargs
5✔
1302
            else:
1303
                args.append(formatted_kwargs)
5✔
1304

1305
        if args:
5✔
1306
            rendered_args = f",\n{self.indentation}".join(args)
5✔
1307
            if len(args) == 1:
5✔
1308
                rendered_args += ","
5✔
1309

1310
            return f"(\n{self.indentation}{rendered_args}\n)"
5✔
1311
        else:
1312
            return ""
5✔
1313

1314
    def render_column_python_type(self, column: Column[Any]) -> str:
5✔
1315
        def get_type_qualifiers() -> tuple[str, TypeEngine[Any], str]:
5✔
1316
            column_type = column.type
5✔
1317
            pre: list[str] = []
5✔
1318
            post_size = 0
5✔
1319
            if column.nullable:
5✔
1320
                self.add_literal_import("typing", "Optional")
5✔
1321
                pre.append("Optional[")
5✔
1322
                post_size += 1
5✔
1323

1324
            if isinstance(column_type, ARRAY):
5✔
1325
                dim = getattr(column_type, "dimensions", None) or 1
5✔
1326
                pre.extend("list[" for _ in range(dim))
5✔
1327
                post_size += dim
5✔
1328

1329
                column_type = column_type.item_type
5✔
1330

1331
            return "".join(pre), column_type, "]" * post_size
5✔
1332

1333
        def render_python_type(column_type: TypeEngine[Any]) -> str:
5✔
1334
            if isinstance(column_type, DOMAIN):
5✔
1335
                column_type = column_type.data_type
5✔
1336

1337
            try:
5✔
1338
                python_type = column_type.python_type
5✔
1339
                python_type_module = python_type.__module__
5✔
1340
                python_type_name = python_type.__name__
5✔
1341
            except NotImplementedError:
5✔
1342
                self.add_literal_import("typing", "Any")
5✔
1343
                return "Any"
5✔
1344

1345
            if python_type_module == "builtins":
5✔
1346
                return python_type_name
5✔
1347

1348
            self.add_module_import(python_type_module)
5✔
1349
            return f"{python_type_module}.{python_type_name}"
5✔
1350

1351
        pre, col_type, post = get_type_qualifiers()
5✔
1352
        column_python_type = f"{pre}{render_python_type(col_type)}{post}"
5✔
1353
        return column_python_type
5✔
1354

1355
    def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
5✔
1356
        column = column_attr.column
5✔
1357
        rendered_column = self.render_column(column, column_attr.name != column.name)
5✔
1358
        rendered_column_python_type = self.render_column_python_type(column)
5✔
1359

1360
        return f"{column_attr.name}: Mapped[{rendered_column_python_type}] = {rendered_column}"
5✔
1361

1362
    def render_relationship(self, relationship: RelationshipAttribute) -> str:
5✔
1363
        def render_column_attrs(column_attrs: list[ColumnAttribute]) -> str:
5✔
1364
            rendered = []
5✔
1365
            for attr in column_attrs:
5✔
1366
                if attr.model is relationship.source:
5✔
1367
                    rendered.append(attr.name)
5✔
1368
                else:
UNCOV
1369
                    rendered.append(repr(f"{attr.model.name}.{attr.name}"))
×
1370

1371
            return "[" + ", ".join(rendered) + "]"
5✔
1372

1373
        def render_foreign_keys(column_attrs: list[ColumnAttribute]) -> str:
5✔
1374
            rendered = []
5✔
1375
            render_as_string = False
5✔
1376
            # Assume that column_attrs are all in relationship.source or none
1377
            for attr in column_attrs:
5✔
1378
                if attr.model is relationship.source:
5✔
1379
                    rendered.append(attr.name)
5✔
1380
                else:
1381
                    rendered.append(f"{attr.model.name}.{attr.name}")
5✔
1382
                    render_as_string = True
5✔
1383

1384
            if render_as_string:
5✔
1385
                return "'[" + ", ".join(rendered) + "]'"
5✔
1386
            else:
1387
                return "[" + ", ".join(rendered) + "]"
5✔
1388

1389
        def render_join(terms: list[JoinType]) -> str:
5✔
1390
            rendered_joins = []
5✔
1391
            for source, source_col, target, target_col in terms:
5✔
1392
                rendered = f"lambda: {source.name}.{source_col} == {target.name}."
5✔
1393
                if target.__class__ is Model:
5✔
1394
                    rendered += "c."
5✔
1395

1396
                rendered += str(target_col)
5✔
1397
                rendered_joins.append(rendered)
5✔
1398

1399
            if len(rendered_joins) > 1:
5✔
UNCOV
1400
                rendered = ", ".join(rendered_joins)
×
UNCOV
1401
                return f"and_({rendered})"
×
1402
            else:
1403
                return rendered_joins[0]
5✔
1404

1405
        # Render keyword arguments
1406
        kwargs: dict[str, Any] = {}
5✔
1407
        if relationship.type is RelationshipType.ONE_TO_ONE and relationship.constraint:
5✔
1408
            if relationship.constraint.referred_table is relationship.source.table:
5✔
1409
                kwargs["uselist"] = False
5✔
1410

1411
        # Add the "secondary" keyword for many-to-many relationships
1412
        if relationship.association_table:
5✔
1413
            table_ref = relationship.association_table.table.name
5✔
1414
            if relationship.association_table.schema:
5✔
1415
                table_ref = f"{relationship.association_table.schema}.{table_ref}"
5✔
1416

1417
            kwargs["secondary"] = repr(table_ref)
5✔
1418

1419
        if relationship.remote_side:
5✔
1420
            kwargs["remote_side"] = render_column_attrs(relationship.remote_side)
5✔
1421

1422
        if relationship.foreign_keys:
5✔
1423
            kwargs["foreign_keys"] = render_foreign_keys(relationship.foreign_keys)
5✔
1424

1425
        if relationship.primaryjoin:
5✔
1426
            kwargs["primaryjoin"] = render_join(relationship.primaryjoin)
5✔
1427

1428
        if relationship.secondaryjoin:
5✔
1429
            kwargs["secondaryjoin"] = render_join(relationship.secondaryjoin)
5✔
1430

1431
        if relationship.backref:
5✔
1432
            kwargs["back_populates"] = repr(relationship.backref.name)
5✔
1433

1434
        rendered_relationship = render_callable(
5✔
1435
            "relationship", repr(relationship.target.name), kwargs=kwargs
1436
        )
1437

1438
        relationship_type: str
1439
        if relationship.type == RelationshipType.ONE_TO_MANY:
5✔
1440
            relationship_type = f"list['{relationship.target.name}']"
5✔
1441
        elif relationship.type in (
5✔
1442
            RelationshipType.ONE_TO_ONE,
1443
            RelationshipType.MANY_TO_ONE,
1444
        ):
1445
            relationship_type = f"'{relationship.target.name}'"
5✔
1446
            if relationship.constraint and any(
5✔
1447
                col.nullable for col in relationship.constraint.columns
1448
            ):
1449
                self.add_literal_import("typing", "Optional")
5✔
1450
                relationship_type = f"Optional[{relationship_type}]"
5✔
1451
        elif relationship.type == RelationshipType.MANY_TO_MANY:
5✔
1452
            relationship_type = f"list['{relationship.target.name}']"
5✔
1453
        else:
UNCOV
1454
            self.add_literal_import("typing", "Any")
×
UNCOV
1455
            relationship_type = "Any"
×
1456

1457
        return (
5✔
1458
            f"{relationship.name}: Mapped[{relationship_type}] "
1459
            f"= {rendered_relationship}"
1460
        )
1461

1462

1463
class DataclassGenerator(DeclarativeGenerator):
5✔
1464
    def __init__(
5✔
1465
        self,
1466
        metadata: MetaData,
1467
        bind: Connection | Engine,
1468
        options: Sequence[str],
1469
        *,
1470
        indentation: str = "    ",
1471
        base_class_name: str = "Base",
1472
        quote_annotations: bool = False,
1473
        metadata_key: str = "sa",
1474
    ):
1475
        super().__init__(
5✔
1476
            metadata,
1477
            bind,
1478
            options,
1479
            indentation=indentation,
1480
            base_class_name=base_class_name,
1481
        )
1482
        self.metadata_key: str = metadata_key
5✔
1483
        self.quote_annotations: bool = quote_annotations
5✔
1484

1485
    def generate_base(self) -> None:
5✔
1486
        self.base = Base(
5✔
1487
            literal_imports=[
1488
                LiteralImport("sqlalchemy.orm", "DeclarativeBase"),
1489
                LiteralImport("sqlalchemy.orm", "MappedAsDataclass"),
1490
            ],
1491
            declarations=[
1492
                (f"class {self.base_class_name}(MappedAsDataclass, DeclarativeBase):"),
1493
                f"{self.indentation}pass",
1494
            ],
1495
            metadata_ref=f"{self.base_class_name}.metadata",
1496
        )
1497

1498

1499
class SQLModelGenerator(DeclarativeGenerator):
5✔
1500
    def __init__(
5✔
1501
        self,
1502
        metadata: MetaData,
1503
        bind: Connection | Engine,
1504
        options: Sequence[str],
1505
        *,
1506
        indentation: str = "    ",
1507
        base_class_name: str = "SQLModel",
1508
    ):
1509
        super().__init__(
5✔
1510
            metadata,
1511
            bind,
1512
            options,
1513
            indentation=indentation,
1514
            base_class_name=base_class_name,
1515
        )
1516

1517
    @property
5✔
1518
    def views_supported(self) -> bool:
5✔
UNCOV
1519
        return False
×
1520

1521
    def render_column_callable(self, is_table: bool, *args: Any, **kwargs: Any) -> str:
5✔
1522
        self.add_import(Column)
5✔
1523
        return render_callable("Column", *args, kwargs=kwargs)
5✔
1524

1525
    def generate_base(self) -> None:
5✔
1526
        self.base = Base(
5✔
1527
            literal_imports=[],
1528
            declarations=[],
1529
            metadata_ref="",
1530
        )
1531

1532
    def collect_imports(self, models: Iterable[Model]) -> None:
5✔
1533
        super(DeclarativeGenerator, self).collect_imports(models)
5✔
1534
        if any(isinstance(model, ModelClass) for model in models):
5✔
1535
            self.remove_literal_import("sqlalchemy", "MetaData")
5✔
1536
            self.add_literal_import("sqlmodel", "SQLModel")
5✔
1537
            self.add_literal_import("sqlmodel", "Field")
5✔
1538

1539
    def collect_imports_for_model(self, model: Model) -> None:
5✔
1540
        super(DeclarativeGenerator, self).collect_imports_for_model(model)
5✔
1541
        if isinstance(model, ModelClass):
5✔
1542
            for column_attr in model.columns:
5✔
1543
                if column_attr.column.nullable:
5✔
1544
                    self.add_literal_import("typing", "Optional")
5✔
1545
                    break
5✔
1546

1547
            if model.relationships:
5✔
1548
                self.add_literal_import("sqlmodel", "Relationship")
5✔
1549

1550
    def render_module_variables(self, models: list[Model]) -> str:
5✔
1551
        declarations: list[str] = []
5✔
1552
        if any(not isinstance(model, ModelClass) for model in models):
5✔
UNCOV
1553
            if self.base.table_metadata_declaration is not None:
×
UNCOV
1554
                declarations.append(self.base.table_metadata_declaration)
×
1555

1556
        return "\n".join(declarations)
5✔
1557

1558
    def render_class_declaration(self, model: ModelClass) -> str:
5✔
1559
        if model.parent_class:
5✔
UNCOV
1560
            parent = model.parent_class.name
×
1561
        else:
1562
            parent = self.base_class_name
5✔
1563

1564
        superclass_part = f"({parent}, table=True)"
5✔
1565
        return f"class {model.name}{superclass_part}:"
5✔
1566

1567
    def render_class_variables(self, model: ModelClass) -> str:
5✔
1568
        variables = []
5✔
1569

1570
        if model.table.name != model.name.lower():
5✔
1571
            variables.append(f"__tablename__ = {model.table.name!r}")
5✔
1572

1573
        # Render constraints and indexes as __table_args__
1574
        table_args = self.render_table_args(model.table)
5✔
1575
        if table_args:
5✔
1576
            variables.append(f"__table_args__ = {table_args}")
5✔
1577

1578
        return "\n".join(variables)
5✔
1579

1580
    def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
5✔
1581
        column = column_attr.column
5✔
1582
        rendered_column = self.render_column(column, True)
5✔
1583
        rendered_column_python_type = self.render_column_python_type(column)
5✔
1584

1585
        kwargs: dict[str, Any] = {}
5✔
1586
        if column.nullable:
5✔
1587
            kwargs["default"] = None
5✔
1588
        kwargs["sa_column"] = f"{rendered_column}"
5✔
1589

1590
        rendered_field = render_callable("Field", kwargs=kwargs)
5✔
1591

1592
        return f"{column_attr.name}: {rendered_column_python_type} = {rendered_field}"
5✔
1593

1594
    def render_relationship(self, relationship: RelationshipAttribute) -> str:
5✔
1595
        rendered = super().render_relationship(relationship).partition(" = ")[2]
5✔
1596
        args = self.render_relationship_args(rendered)
5✔
1597
        kwargs: dict[str, Any] = {}
5✔
1598
        annotation = repr(relationship.target.name)
5✔
1599

1600
        if relationship.type in (
5✔
1601
            RelationshipType.ONE_TO_MANY,
1602
            RelationshipType.MANY_TO_MANY,
1603
        ):
1604
            annotation = f"list[{annotation}]"
5✔
1605
        else:
1606
            self.add_literal_import("typing", "Optional")
5✔
1607
            annotation = f"Optional[{annotation}]"
5✔
1608

1609
        rendered_field = render_callable("Relationship", *args, kwargs=kwargs)
5✔
1610
        return f"{relationship.name}: {annotation} = {rendered_field}"
5✔
1611

1612
    def render_relationship_args(self, arguments: str) -> list[str]:
5✔
1613
        argument_list = arguments.split(",")
5✔
1614
        # delete ')' and ' ' from args
1615
        argument_list[-1] = argument_list[-1][:-1]
5✔
1616
        argument_list = [argument[1:] for argument in argument_list]
5✔
1617

1618
        rendered_args: list[str] = []
5✔
1619
        for arg in argument_list:
5✔
1620
            if "back_populates" in arg:
5✔
1621
                rendered_args.append(arg)
5✔
1622
            if "uselist=False" in arg:
5✔
1623
                rendered_args.append("sa_relationship_kwargs={'uselist': False}")
5✔
1624

1625
        return rendered_args
5✔
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc