src/baby_monitor/models/db/user.py

@@ -1,9 +1,14 @@
 """Database models."""
 
+from typing import TYPE_CHECKING
 from sqlalchemy import Column, Integer, String, DateTime, Boolean
 from datetime import datetime
 
-from baby_monitor.repositories.dependencies.get_database import Base
+if TYPE_CHECKING:
+    from sqlalchemy.orm import DeclarativeBase
+    Base = DeclarativeBase
+else:
+    from baby_monitor.repositories.dependencies.get_database import Base
 
 
 class User(Base):

src/baby_monitor/models/invitation.py

@@ -5,10 +5,15 @@ SQLite will store as naive UTC, which is handled in the repository layer.
 """
 
 from datetime import datetime
+from typing import TYPE_CHECKING
 from sqlalchemy import Boolean, DateTime, Integer, String
 from sqlalchemy.orm import Mapped, mapped_column
 
-from baby_monitor.repositories.dependencies.get_database import Base
+if TYPE_CHECKING:
+    from sqlalchemy.orm import DeclarativeBase
+    Base = DeclarativeBase
+else:
+    from baby_monitor.repositories.dependencies.get_database import Base
 
 
 class Invitation(Base):

src/baby_monitor/repositories/dependencies/get_database.py

@@ -3,9 +3,13 @@
 import os
 from pathlib import Path
 from collections.abc import Generator
+from typing import TYPE_CHECKING
 from sqlalchemy import create_engine
 from sqlalchemy.orm import declarative_base, sessionmaker, Session
 
+if TYPE_CHECKING:
+    from sqlalchemy.orm import DeclarativeBase
+
 # Database directory
 DATA_DIR = Path(os.getenv("DATA_DIR", "/data"))
 
@@ -23,7 +27,7 @@ engine = create_engine(
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 # Base class for declarative models
-Base = declarative_base()
+Base: "type[DeclarativeBase]" = declarative_base()
 
 
 def init_db() -> None: