fix circular imports

affine/engine/__init__.py

@@ -1,8 +1,9 @@
+from .base import Engine
 from .local import LocalEngine
 
 try:
     from .qdrant import QdrantEngine
 except ModuleNotFoundError:
     pass
 
-__all__ = ["LocalEngine", "QdrantEngine"]
+__all__ = ["Engine", "LocalEngine", "QdrantEngine"]

affine/engine/base.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+
+from affine.collection import Collection, FilterSet
+
+
+class Engine(ABC):
+    @abstractmethod
+    def query(self, filter_set: FilterSet) -> list[Collection]:
+        pass
+
+    @abstractmethod
+    def insert(self, record: Collection) -> int:
+        pass
+
+    @abstractmethod
+    def delete(self, collection: type, id_: int) -> None:
+        pass
+
+    @abstractmethod
+    def get_elements_by_ids(
+        self, collection: type, ids: list[int]
+    ) -> list[Collection]:
+        pass
+
+    def get_element_by_id(self, collection: type, id_: int) -> Collection:
+        ret = self.get_elements_by_ids(collection, [id_])
+        if len(ret) == 0:
+            raise ValueError(f"No record found with id {id_}")
+        if len(ret) > 1:
+            raise ValueError(f"Multiple records found with id {id_}")
+        return ret[0]

affine/engine/local.py

@@ -1,12 +1,12 @@
 import pickle
-from abc import ABC, abstractmethod
 from collections import defaultdict
 from pathlib import Path
 from typing import BinaryIO
 
 import numpy as np
 
 from affine.collection import Collection, Filter, FilterSet
+from affine.engine import Engine
 
 
 def apply_filter_to_record(filter_: Filter, record: Collection) -> bool:
@@ -72,34 +72,6 @@ def apply_filters_to_records(
     return records
 
 
-class Engine(ABC):
-    @abstractmethod
-    def query(self, filter_set: FilterSet) -> list[Collection]:
-        pass
-
-    @abstractmethod
-    def insert(self, record: Collection) -> int:
-        pass
-
-    @abstractmethod
-    def delete(self, collection: type, id_: int) -> None:
-        pass
-
-    @abstractmethod
-    def get_elements_by_ids(
-        self, collection: type, ids: list[int]
-    ) -> list[Collection]:
-        pass
-
-    def get_element_by_id(self, collection: type, id_: int) -> Collection:
-        ret = self.get_elements_by_ids(collection, [id_])
-        if len(ret) == 0:
-            raise ValueError(f"No record found with id {id_}")
-        if len(ret) > 1:
-            raise ValueError(f"Multiple records found with id {id_}")
-        return ret[0]
-
-
 class LocalEngine(Engine):
     def __init__(self) -> None:  # maybe add option to the init for ANN algo
         self.records: dict[str, list[Collection]] = defaultdict(list)