Remove if TYPE_CHECKING clause from most model files

lightly_studio/src/lightly_studio/__init__.py

@@ -1,6 +1,10 @@
 # Set up logging before importing any other modules.
 # Add noqa to silence unused import and unsorted imports linter warnings.
 from . import setup_logging  # noqa: F401 I001
+
+# Import db_manager for SQLModel to discover db models.
+from lightly_studio import db_manager  # noqa: F401
+
 from lightly_studio.core.image_dataset import ImageDataset as Dataset
 from lightly_studio.core.start_gui import start_gui
 from lightly_studio.models.collection import SampleType

lightly_studio/src/lightly_studio/api/db_tables.py

@@ -1,5 +1,9 @@
 """Module provides functions to initialize and manage the DuckDB."""
 
+# Import SampleTable first to avoid circular imports. Needs I001 to prevent import sorting.
+from lightly_studio.models.sample import (  # noqa: I001
+    SampleTable,  # noqa: F401, required for SQLModel to work properly
+)
 from lightly_studio.models.annotation.annotation_base import (
     AnnotationBaseTable,  # noqa: F401, required for SQLModel to work properly
 )
@@ -18,9 +22,6 @@
 from lightly_studio.models.metadata import (
     SampleMetadataTable,  # noqa: F401, required for SQLModel to work properly
 )
-from lightly_studio.models.sample import (
-    SampleTable,  # noqa: F401, required for SQLModel to work properly
-)
 from lightly_studio.models.sample_embedding import (
     SampleEmbeddingTable,  # noqa: F401, required for SQLModel to work properly
 )

lightly_studio/src/lightly_studio/models/annotation_label.py

@@ -1,19 +1,15 @@
 """This module defines the AnnotationLabel model for the application."""
 
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, List
+from typing import List
 from uuid import UUID, uuid4
 
 from sqlalchemy.orm import Mapped
 from sqlmodel import Field, Relationship, SQLModel, UniqueConstraint
 
-if TYPE_CHECKING:
-    from lightly_studio.models.annotation.annotation_base import (
-        AnnotationBaseTable,
-    )
-
-else:
-    AnnotationBaseTable = object
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
 
 
 class AnnotationLabelBase(SQLModel):

lightly_studio/src/lightly_studio/models/metadata.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel
@@ -16,12 +16,7 @@
     deserialize_complex_metadata,
     serialize_complex_metadata,
 )
-
-if TYPE_CHECKING:
-    from lightly_studio.models.sample import SampleTable
-else:
-    SampleTable = object
-
+from lightly_studio.models.sample import SampleTable
 
 TYPE_TO_NAME_MAP = {
     bool: "boolean",

lightly_studio/src/lightly_studio/models/sample.py

@@ -1,34 +1,14 @@
 """This module defines the Sample model for the application."""
 
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import Any, List, Optional
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, ConfigDict
 from pydantic import Field as PydanticField
 from sqlalchemy.orm import Mapped, Session
 from sqlmodel import Field, Relationship, SQLModel
 
-from lightly_studio.models.metadata import (
-    SampleMetadataTable,
-    SampleMetadataView,
-)
-from lightly_studio.resolvers import metadata_resolver
-
-if TYPE_CHECKING:
-    from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable, AnnotationView
-    from lightly_studio.models.caption import CaptionTable, CaptionView
-    from lightly_studio.models.sample_embedding import SampleEmbeddingTable
-    from lightly_studio.models.tag import TagTable
-
-else:
-    AnnotationBaseTable = object
-    TagTable = object
-    SampleEmbeddingTable = object
-    CaptionTable = object
-    CaptionView = object
-    AnnotationView = object
-
 
 class SampleTagLinkTable(SQLModel, table=True):
     """Model to define links between Sample and Tag Many-to-Many."""
@@ -115,6 +95,9 @@ def __setitem__(self, key: str, value: Any) -> None:
         if session is None:
             raise RuntimeError("No database session found for this instance")
 
+        # Delayed import to avoid circular dependencies.
+        from lightly_studio.resolvers import metadata_resolver
+
         # Use metadata_resolver to handle the database operations.
         # Added type: ignore to avoid type checking issues. SQLAlchemy and
         # SQLModel sessions are compatible at runtime but have different type
@@ -136,7 +119,7 @@ class SampleView(SampleBase):
 
     tags: List["TagTable"] = []
     metadata_dict: Optional["SampleMetadataView"] = None
-    captions: List[CaptionView] = []
+    captions: List["CaptionView"] = []
     annotations: List["AnnotationView"] = []
 
 
@@ -148,3 +131,20 @@ class SampleViewsWithCount(BaseModel):
     samples: List[SampleView] = PydanticField(..., alias="data")
     total_count: int
     next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")
+
+
+# Import at the bottom to:
+# 1) avoid circular imports
+# 2) satisfy mypy
+# 3) include types in schema generation
+from lightly_studio.models.annotation.annotation_base import (  # noqa: E402
+    AnnotationBaseTable,
+    AnnotationView,
+)
+from lightly_studio.models.caption import CaptionTable, CaptionView  # noqa: E402
+from lightly_studio.models.metadata import (  # noqa: E402
+    SampleMetadataTable,
+    SampleMetadataView,
+)
+from lightly_studio.models.sample_embedding import SampleEmbeddingTable  # noqa: E402
+from lightly_studio.models.tag import TagTable  # noqa: E402

lightly_studio/src/lightly_studio/models/sample_embedding.py

@@ -2,17 +2,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 from uuid import UUID
 
 from sqlalchemy import ARRAY, Float
 from sqlmodel import Column, Field, Relationship, SQLModel
 
-if TYPE_CHECKING:
-    from lightly_studio.models.sample import SampleTable
-
-else:
-    SampleTable = object
+from lightly_studio.models.sample import SampleTable
 
 
 class SampleEmbeddingBase(SQLModel):

lightly_studio/src/lightly_studio/models/tag.py

@@ -1,27 +1,18 @@
 """This module contains the Tag model and related enumerations."""
 
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, List, Literal, Optional
+from typing import List, Literal, Optional
 from uuid import UUID, uuid4
 
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.orm import Mapped
 from sqlmodel import Field, Relationship, SQLModel, String
 
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
 from lightly_studio.models.annotation.links import AnnotationTagLinkTable
-from lightly_studio.models.sample import SampleTagLinkTable
-
-if TYPE_CHECKING:
-    from lightly_studio.models.annotation.annotation_base import (
-        AnnotationBaseTable,
-    )
-    from lightly_studio.models.sample import SampleTable
-
-else:
-    SampleTable = object
-    TagTable = object
-    AnnotationBaseTable = object
-
+from lightly_studio.models.sample import SampleTable, SampleTagLinkTable
 
 # TagKind is the kind of tag we support.
 TagKind = Literal[

lightly_studio/src/lightly_studio/models/video.py

@@ -1,6 +1,6 @@
 """This module defines the Video and VideoFrame model for the application."""
 
-from typing import TYPE_CHECKING, List, Optional
+from typing import List, Optional
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict
@@ -11,13 +11,6 @@
 from lightly_studio.models.range import FloatRange, IntRange
 from lightly_studio.models.sample import SampleTable, SampleView
 
-if TYPE_CHECKING:
-    from lightly_studio.models.annotation.annotation_base import (
-        AnnotationBaseTable,
-    )
-else:
-    AnnotationBaseTable = object
-
 
 class VideoBase(SQLModel):
     """Base class for the Video model."""

lightly_studio_view/src/lib/schema.d.ts

@@ -2738,18 +2738,18 @@ export interface components {
              * Tags
              * @default []
              */
-            tags: unknown[];
+            tags: components["schemas"]["TagTable"][];
             metadata_dict?: components["schemas"]["SampleMetadataView"] | null;
             /**
              * Captions
              * @default []
              */
-            captions: unknown[];
+            captions: components["schemas"]["CaptionView"][];
             /**
              * Annotations
              * @default []
              */
-            annotations: unknown[];
+            annotations: components["schemas"]["AnnotationView"][];
         };
         /**
          * SampleViewsWithCount