make save and load os independent; add uv to pre-commit to .toml

.pre-commit-config.yaml

@@ -12,4 +12,9 @@ repos:
     - id: ruff-format
       types_or: [ python, pyi ]
       files: .*
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    # uv version.
+    rev: 0.6.10
+    hooks:
+      - id: uv-lock
 files: ''

baguetter/indices/dense/faiss.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Any
 
@@ -15,7 +16,6 @@
     from collections.abc import Callable
 
     from baguetter.types import DTypeLike, Key, ScoringMetric, TextOrVector
-    from baguetter.utils.file_repository import AbstractFileRepository
 
 
 def _support_nprobe(index: Any) -> bool:
@@ -93,7 +93,6 @@ def from_config(cls, config: FaissDenseIndexConfig) -> FaissDenseIndex:
     def _save(
         self,
         path: str,
-        repository: AbstractFileRepository,
     ) -> str:
         """Save the index state and data."""
         state = {
@@ -106,10 +105,10 @@ def _save(
             index_file_path,
         ) = BaseDenseIndex.build_index_file_paths(path or self.name)
 
-        with repository.open(state_file_path, "wb") as file:
+        with open(state_file_path, "wb") as file:
             np.savez_compressed(file, state=state)
 
-        with repository.open(index_file_path, "wb") as file:
+        with open(index_file_path, "wb") as file:
             index = faiss.serialize_index(self.faiss_index)
             np.savez_compressed(file, index=index)
         return state_file_path
@@ -119,7 +118,6 @@ def _load(
         cls,
         path: str,
         *,
-        repository: AbstractFileRepository,
         mmap: bool = False,
     ) -> FaissDenseIndex:
         """Load the index from saved state.
@@ -137,22 +135,22 @@ def _load(
         """
         state_file_path, index_file_path = BaseDenseIndex.build_index_file_paths(path)
 
-        if not repository.exists(state_file_path):
+        if not os.path.exists(state_file_path):
             msg = f"Index.state {state_file_path} not found in repository."
             raise FileNotFoundError(msg)
 
-        if not repository.exists(index_file_path):
+        if not os.path.exists(index_file_path):
             msg = f"Index.index {index_file_path} not found in repository."
             raise FileNotFoundError(msg)
 
-        with repository.open(state_file_path, "rb") as file:
+        with open(state_file_path, "rb") as file:
             state = np.load(file, allow_pickle=True)["state"][()]
             config = state["config"]
             index = cls.from_config(FaissDenseIndexConfig(**config))
             index.key_mapping = state["key_mapping"]
 
         mmap_mode = "r" if mmap else None
-        with repository.open(index_file_path, "rb") as file:
+        with open(index_file_path, "rb") as file:
             stored = np.load(file, allow_pickle=True, mmap_mode=mmap_mode)
             index.faiss_index = faiss.deserialize_index(stored["index"])
 

baguetter/indices/dense/usearch.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import math
 import tempfile
 from dataclasses import asdict
@@ -17,7 +18,6 @@
     from collections.abc import Callable
 
     from baguetter.types import DTypeLike, Key, ScoringMetric, TextOrVector
-    from baguetter.utils.file_repository import AbstractFileRepository
 
 
 def get_normalization_fn(metric: MetricKind, embedding_dim: int) -> Callable[[np.ndarray], np.ndarray]:
@@ -141,7 +141,6 @@ def from_config(cls, config: UsearchDenseIndexConfig) -> USearchDenseIndex:
     def _save(
         self,
         path: str,
-        repository: AbstractFileRepository,
     ) -> str:
         """Save the index state and data.
 
@@ -161,12 +160,12 @@ def _save(
             index_file_path,
         ) = BaseDenseIndex.build_index_file_paths(path or self.name)
 
-        with repository.open(state_file_path, "wb") as file:
+        with open(state_file_path, "wb") as file:
             np.savez_compressed(file, state=state)
 
         with (
-            repository.open(index_file_path, "wb") as file,
-            tempfile.NamedTemporaryFile() as temp_file,
+            open(index_file_path, "wb") as file,
+            tempfile.NamedTemporaryFile(delete=False) as temp_file,
         ):
             Index.save(self, temp_file.name)
             file.write(temp_file.read())
@@ -178,7 +177,6 @@ def _load(
         cls,
         path: str,
         *,
-        repository: AbstractFileRepository,
         mmap: bool = False,
     ) -> USearchDenseIndex:
         """Load the index from saved state.
@@ -200,24 +198,24 @@ def _load(
             index_file_path,
         ) = BaseDenseIndex.build_index_file_paths(path)
 
-        if not repository.exists(state_file_path):
+        if not os.path.exists(state_file_path):
             msg = f"Index.state {state_file_path} not found in repository."
             raise FileNotFoundError(msg)
 
-        if not repository.exists(index_file_path):
+        if not os.path.exists(index_file_path):
             msg = f"Index.index {index_file_path} not found in repository."
             raise FileNotFoundError(msg)
 
-        with repository.open(state_file_path, "rb") as file:
+        with open(state_file_path, "rb") as file:
             state = np.load(file, allow_pickle=True)["state"][()]
             config = state["config"]
             index = cls.from_config(UsearchDenseIndexConfig(**config))
             index.key_mapping = state["key_mapping"]
             index.key_counter = state["id_count"]
 
         with (
-            repository.open(index_file_path, "rb") as file,
-            tempfile.NamedTemporaryFile(delete=not mmap) as temp,
+            open(index_file_path, "rb") as file,
+            tempfile.NamedTemporaryFile(delete=False) as temp,
         ):
             temp.write(file.read())
             temp.seek(0)

baguetter/indices/sparse/base.py

@@ -19,7 +19,6 @@
     from collections.abc import Generator
 
     from baguetter.types import Key, TextOrTokens
-    from baguetter.utils.file_repository import AbstractFileRepository
 
 
 class BaseSparseIndex(BaseIndex, abc.ABC):
@@ -162,7 +161,6 @@ def vocabulary(self) -> dict[str, int]:
     def _save(
         self,
         path: str,
-        repository: AbstractFileRepository,
     ) -> str:
         """Save the index to the given path.
 
@@ -177,15 +175,14 @@ def _save(
             "corpus_tokens": self.corpus_tokens,
             "config": dataclasses.asdict(self.config),
         }
-        with repository.open(path, "wb") as f:
+        with open(path, "wb") as f:
             np.savez_compressed(f, state=state)
         return path
 
     @classmethod
     def _load(
         cls,
         path: str,
-        repository: AbstractFileRepository,
         *,
         mmap: bool = False,
     ) -> BaseSparseIndex:
@@ -203,12 +200,12 @@ def _load(
             FileNotFoundError: If the index file is not found.
 
         """
-        if not repository.exists(path):
+        if not os.path.exists(path):
             msg = f"Index {path} not found."
             raise FileNotFoundError(msg)
 
         mmap_mode = "r" if mmap else None
-        with repository.open(path, "rb") as f:
+        with open(path, "rb") as f:
             stored = np.load(f, allow_pickle=True, mmap_mode=mmap_mode)
             state = stored["state"][()]
             retriever = cls.from_config(SparseIndexConfig(**state["config"]))

baguetter/utils/file_repository.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from pathlib import Path
 
 from fsspec import AbstractFileSystem
@@ -103,12 +104,13 @@ def __init__(self, path: str | None = None, **kwargs) -> None:
 
         """
         super().__init__(**kwargs)
-        self._base_path = str(Path(path or f"{settings.base_path}/repository").resolve())
-        if not self.isdir(self._base_path):
-            if self.exists(self._base_path):
-                msg = f"Path '{self._base_path}' exists but is not a directory."
+        base_path = str(Path(path or f"{settings.base_path}/repository").resolve())
+        if not os.path.isdir(base_path):
+            if os.path.exists(base_path):
+                msg = f"Path '{base_path}' exists but is not a directory."
                 raise ValueError(msg)
-            self.mkdir(self._base_path, parents=True, exist_ok=True)
+            os.makedirs(base_path, exist_ok=True)
+        self._base_path = base_path
 
     def _strip_protocol(self, path: str) -> str:
         """Strip the protocol from the given path.

baguetter/utils/persistable.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import os
 from abc import ABC, abstractmethod
 from typing import Any
 
-from baguetter.utils.file_repository import AbstractFileRepository, HuggingFaceFileRepository, LocalFileRepository
-
+from baguetter.utils.file_repository import HuggingFaceFileRepository
 
 class Persistable(ABC):
     """Abstract base class for objects that can be persisted to and loaded from storage.
@@ -18,7 +18,6 @@ class Persistable(ABC):
     def _load(
         cls,
         path: str,
-        repository: AbstractFileRepository,
         *,
         allow_pickle: bool = True,
         mmap: bool = False,
@@ -27,7 +26,6 @@ def _load(
 
         Args:
             path (str): Path of the object to load.
-            repository (AbstractFileRepository): File repository to load from.
             allow_pickle (bool, optional): Whether to allow loading pickled objects. Defaults to True.
             mmap (bool, optional): Whether to memory-map the file. Defaults to False.
 
@@ -37,12 +35,11 @@ def _load(
         """
 
     @abstractmethod
-    def _save(self, path: str, repository: AbstractFileRepository) -> str:
+    def _save(self, path: str) -> str:
         """Save the object to storage.
 
         Args:
             path (str): Path to save the object to.
-            repository (AbstractFileRepository): File repository to save to.
 
         Returns:
             str: Path to the saved object.
@@ -59,12 +56,8 @@ def save(self, path: str) -> str:
             str: Path to the saved object.
 
         """
-        repository = LocalFileRepository()
-        directory = path.rsplit("/", 1)
-        if len(directory) > 1:
-            repository.mkdirs(directory[0], exist_ok=True)
-        path = self._save(path=path, repository=repository)
-        return repository.info(path)["name"]
+        path = self._save(path=path)
+        return os.path.basename(path)
 
     @classmethod
     def load(cls, path: str, *, mmap: bool = False) -> Any:
@@ -78,10 +71,8 @@ def load(cls, path: str, *, mmap: bool = False) -> Any:
             Any: The loaded object.
 
         """
-        repository = LocalFileRepository()
         return cls._load(
             path=path,
-            repository=repository,
             mmap=mmap,
         )
 

pyproject.toml

@@ -52,3 +52,8 @@ include = ["baguetter", "baguetter.*"]
 
 [tool.setuptools.package-data]
 baguetter = ["py.typed"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.5",
+]

tests/indices/dense/faiss_test.py

@@ -6,7 +6,6 @@
 
 from baguetter.indices.base import SearchResults
 from baguetter.indices.dense.faiss import FaissDenseIndex
-from baguetter.utils.file_repository import LocalFileRepository
 
 
 @pytest.fixture
@@ -100,10 +99,9 @@ def test_faiss_save_and_load(sample_data):
 
     with tempfile.TemporaryDirectory() as tmpdir:
         save_path = "faiss_index"
-        repo = LocalFileRepository(tmpdir)
-        index._save(repository=repo, path=save_path)
+        index._save(path=save_path)
 
-        loaded_index = FaissDenseIndex._load(save_path, repository=repo)
+        loaded_index = FaissDenseIndex._load(save_path)
 
         assert loaded_index.size == index.size
         assert loaded_index.config.__dict__ == index.config.__dict__

tests/indices/dense/usearch_test.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import tempfile
 
 import numpy as np
@@ -8,7 +9,6 @@
 from baguetter.indices.dense.base import _INDEX_PREFIX
 from baguetter.indices.dense.config import DenseIndexConfig
 from baguetter.indices.dense.usearch import USearchDenseIndex
-from baguetter.utils.file_repository import LocalFileRepository
 
 
 @pytest.fixture
@@ -94,13 +94,12 @@ def test_usearch_save_load(sample_data):
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         save_path = "usearch_new-index"
-        repository = LocalFileRepository(tmp_dir)
 
         # Save the index
-        index._save(path=save_path, repository=repository)
+        index._save(path=save_path)
 
         # Load the index
-        loaded_index = USearchDenseIndex._load(path=save_path, repository=repository)
+        loaded_index = USearchDenseIndex._load(path=save_path)
 
         assert loaded_index.config.embedding_dim == 128
         assert loaded_index.key_counter == 10
@@ -116,8 +115,8 @@ def test_usearch_save_load(sample_data):
         )
 
         # Verify file names
-        assert repository.exists(save_path)
-        assert repository.exists(f"{_INDEX_PREFIX}{save_path}")
+        assert os.path.exists(save_path)
+        assert os.path.exists(f"{_INDEX_PREFIX}{save_path}")
 
 
 def test_usearch_embed_function():