diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index fe37cef..6049718 100644 --- a/src/sqlacodegen/generators.py +++ b/src/sqlacodegen/generators.py @@ -1385,6 +1385,7 @@ def collect_imports(self, models: Iterable[Model]) -> None: super(DeclarativeGenerator, self).collect_imports(models) if any(isinstance(model, Model) for model in models): self.add_literal_import("sqlalchemy", "MetaData") + if any(isinstance(model, ModelClass) for model in models): self.add_literal_import("sqlmodel", "SQLModel") self.add_literal_import("sqlmodel", "Field")