Type data store for embedders properly

Renumics · Dec 1, 2023 · a70ecf4 · a70ecf4
1 parent 2580fd6
commit a70ecf4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/renumics/spotlight/embeddings/__init__.py b/renumics/spotlight/embeddings/__init__.py
@@ -4,9 +4,10 @@
 
 import importlib
 import pkgutil
-from typing import Any, Dict, List
+from typing import Dict, List
 
 import numpy as np
+from renumics.spotlight.data_store import DataStore
 
 from renumics.spotlight.embeddings.typing import Embedder
 from .registry import registered_embedders
@@ -17,7 +18,7 @@
     importlib.import_module(embedders_namespace.__name__ + "." + module_info.name)
 
 
-def create_embedders(data_store: Any, columns: List[str]) -> Dict[str, Embedder]:
+def create_embedders(data_store: DataStore, columns: List[str]) -> Dict[str, Embedder]:
     """
     Create embedding functions for the given data store.
     """

diff --git a/renumics/spotlight/embeddings/typing.py b/renumics/spotlight/embeddings/typing.py
@@ -4,16 +4,18 @@
 
 import itertools
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Iterable, List, Optional
+from typing import Callable, Iterable, List, Optional
 
 import numpy as np
 
+from renumics.spotlight.data_store import DataStore
+
 
 class Embedder(ABC):
-    data_store: Any
+    data_store: DataStore
     column: str
 
-    def __init__(self, data_store: Any, column: str) -> None:
+    def __init__(self, data_store: DataStore, column: str) -> None:
         self.data_store = data_store
         self.column = column
 
@@ -36,7 +38,7 @@ class FunctionalEmbedder(Embedder):
 
     def __init__(
         self,
-        data_store: Any,
+        data_store: DataStore,
         column: str,
         preprocess_func: PreprocessFunc,
         embed_func: EmbedFunc,