refactor: use type[_]

libs/kotaemon/kotaemon/deps/container.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from functools import cached_property, lru_cache
 from pathlib import Path
-from typing import Any, Sequence
+from typing import Any, Sequence, TypeAlias
 
 from ktem.index.file.base import BaseFileIndexIndexing, BaseFileIndexRetriever
 from ktem.index.file.pipelines import DocumentRetrievalPipeline, IndexDocumentPipeline
@@ -19,6 +19,8 @@
 from ..storages.vectorstores import BaseVectorStore, ChromaVectorStore
 from .registry import Dependency, Registry
 
+SourceRecords: TypeAlias = Sequence[Row[tuple[Source]]]
+
 
 @lru_cache(1)
 def get_engine() -> Engine:
@@ -58,19 +60,19 @@ def __init__(self, collection_idx: int, private: bool = False):
         self.private = private
 
     @cached_property
-    def source(self) -> Source:
+    def source(self) -> type[Source]:
         source = Source.from_index(self.collection_idx, self.private)
         source.metadata.create_all(self.engine)
         return source
 
     @cached_property
-    def index(self) -> Index:
+    def index(self) -> type[Index]:
         index = Index.from_index(self.collection_idx)
         index.metadata.create_all(self.engine)
         return index
 
     @cached_property
-    def filegroup(self) -> FileGroup:
+    def filegroup(self) -> type[FileGroup]:
         filegroup = FileGroup.from_index(self.collection_idx)
         filegroup.metadata.create_all(self.engine)
         return filegroup
@@ -87,11 +89,12 @@ def filestorage_path(self) -> Path:
 class FileCRUD:
     engine = get_engine()
 
+    def __init__(self, source: type[Source]):
+        self.source = source
+
     def list_docids(self) -> list[str]:
         with Session(self.engine) as session:
-            records: Sequence[Row[tuple[Source]]] = session.execute(
-                select(Source)
-            ).all()
+            records: SourceRecords = session.execute(select(self.source)).all()
             return [record[0].id for record in records]
 
 
@@ -179,7 +182,7 @@ class Container:
     user_idx: int = 1
     private: bool = False
     fileschema: FileSchemaFactory = FileSchemaFactory(collection_idx, private)
-    crud: FileCRUD = FileCRUD()
+    crud: FileCRUD = FileCRUD(fileschema.source)
 
     vectorstores: Registry[BaseVectorStore] = Registry(
         {"chroma": Dependency(VectorstoreFactory.chroma)}

libs/kotaemon/kotaemon/schemas/file/base.py

@@ -7,6 +7,6 @@ class BaseSchema(DeclarativeBase):
     __abstract__ = True
 
     @classmethod
-    def from_dict(cls, cls_name: str, params: dict[str, Any]) -> Any:
+    def from_dict(cls, cls_name: str, params: dict[str, Any]) -> type["BaseSchema"]:
         params["__abstract__"] = False
         return type(cls_name, (BaseSchema,), params)

libs/kotaemon/kotaemon/schemas/file/group.py

@@ -2,7 +2,6 @@
 
 from sqlalchemy import JSON, Column, DateTime, Integer, String, func
 from sqlalchemy.ext.mutable import MutableDict
-from typing_extensions import Self
 
 from .base import BaseSchema
 
@@ -17,6 +16,6 @@ class FileGroup(BaseSchema):
     data = Column(MutableDict.as_mutable(JSON), default={"files": []})  # type: ignore
 
     @classmethod
-    def from_index(cls, idx: int) -> Self:
+    def from_index(cls, idx: int) -> type["FileGroup"]:
         cls.__tablename__ = f"index__{idx}__group"
-        return cast(Self, cls.from_dict("FileGroup", dict(vars(cls))))
+        return cast(type["FileGroup"], cls.from_dict("FileGroup", dict(vars(cls))))

libs/kotaemon/kotaemon/schemas/file/index.py

@@ -1,7 +1,6 @@
 from typing import cast
 
 from sqlalchemy import Column, Integer, String
-from typing_extensions import Self
 
 from .base import BaseSchema
 
@@ -16,6 +15,6 @@ class Index(BaseSchema):
     user = Column(Integer, default=1)
 
     @classmethod
-    def from_index(cls, idx: int) -> Self:
+    def from_index(cls, idx: int) -> type["Index"]:
         cls.__tablename__ = f"index__{idx}__index"
-        return cast(Self, cls.from_dict("Index", dict(vars(cls))))
+        return cast(type["Index"], cls.from_dict("Index", dict(vars(cls))))

libs/kotaemon/kotaemon/schemas/file/source.py

@@ -3,7 +3,6 @@
 
 from sqlalchemy import JSON, Column, DateTime, Integer, String, UniqueConstraint, func
 from sqlalchemy.ext.mutable import MutableDict
-from typing_extensions import Self
 
 from .base import BaseSchema
 
@@ -22,7 +21,7 @@ class Source(BaseSchema):
     note = Column(MutableDict.as_mutable(JSON), default={})  # type: ignore
 
     @classmethod
-    def from_index(cls, idx: int, private: bool = False) -> Self:
+    def from_index(cls, idx: int, private: bool = False) -> type["Source"]:
         cls.__tablename__ = f"index__{idx}__source"
         if private:
             cls.__table_args__ = (
@@ -32,4 +31,4 @@ def from_index(cls, idx: int, private: bool = False) -> Self:
         else:
             cls.name = Column(String, unique=True)
 
-        return cast(Self, cls.from_dict("Source", dict(vars(cls))))
+        return cast(type["Source"], cls.from_dict("Source", dict(vars(cls))))