Repository: chr1st1ank/narrow-down (language: python)

Issue:
Small interface refinements
- The initialize() function could be called by SimilarityStore and it could be handled in a way, that it doesn't matter if it is called multiple times.
- The data_types module has a too generic name
- 


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ad70c62..9ca3729 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - ScyllaDBStore now accepts a `table_prefix` setting.
 
+### Changed
+- The classes in narrow_down.data_types were moved to narrow_down.storage.
+- The `initialize()` method of the storage backends can now be called multiple times without issues.
+
 ### Fixed
 - A use of collections.Counter as typehint broke mypy checks.  
 
diff --git a/docs/user_guide/basic_usage.ipynb b/docs/user_guide/basic_usage.ipynb
index 9fea5e9..9893b9d 100644
--- a/docs/user_guide/basic_usage.ipynb
+++ b/docs/user_guide/basic_usage.ipynb
@@ -14,21 +14,21 @@
     "## Indexing\n",
     "The code block below shows how to create and configure a SimilarityStore() object. \n",
     "\n",
-    "Here we choose the [StorageLevel](narrow_down.data_types.StorageLevel) `Document`, so that the whole document is stored in the `SimilarityStore` and can be returned from it. A target similarity threshold of 75% is defined, which means that we want to search for documents which have a Jaccard similarity of at least 75% with the input document. To calculate the similarity, the document is first preprocessed by a tokenizer function. Here we choose the character 3-grams of a document for this."
+    "Here we choose the [StorageLevel](narrow_down.storage.StorageLevel) `Document`, so that the whole document is stored in the `SimilarityStore` and can be returned from it. A target similarity threshold of 75% is defined, which means that we want to search for documents which have a Jaccard similarity of at least 75% with the input document. To calculate the similarity, the document is first preprocessed by a tokenizer function. Here we choose the character 3-grams of a document for this."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "id": "637d11ab",
    "metadata": {},
    "outputs": [],
    "source": [
     "import narrow_down as nd\n",
-    "from narrow_down.data_types import StoredDocument\n",
+    "from narrow_down.storage import StorageLevel, StoredDocument\n",
     "\n",
     "similarity_store = await nd.similarity_store.SimilarityStore.create(\n",
-    "    storage_level=nd.data_types.StorageLevel.Document,\n",
+    "    storage_level=StorageLevel.Document,\n",
     "    similarity_threshold=0.75,\n",
     "    tokenize=\"char_ngrams(3)\",\n",
     ")"
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "62713b60",
    "metadata": {},
    "outputs": [],
@@ -100,7 +100,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "id": "c3eca9d6",
    "metadata": {},
    "outputs": [
@@ -110,7 +110,7 @@
        "True"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -122,7 +122,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "146850e2",
    "metadata": {},
    "outputs": [
@@ -132,7 +132,7 @@
        "True"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -144,7 +144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "89976533",
    "metadata": {},
    "outputs": [
@@ -154,7 +154,7 @@
        "True"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -166,7 +166,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "35dc5cf1",
    "metadata": {},
    "outputs": [
@@ -176,7 +176,7 @@
        "True"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -198,7 +198,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "id": "7d911c31",
    "metadata": {},
    "outputs": [
@@ -208,7 +208,7 @@
        "42"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -216,6 +216,14 @@
    "source": [
     "await similarity_store.insert(\"Good cookie\", document_id=42)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3126d3a8",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/docs/user_guide/storage_backends.ipynb b/docs/user_guide/storage_backends.ipynb
index 83eb024..c70ca97 100644
--- a/docs/user_guide/storage_backends.ipynb
+++ b/docs/user_guide/storage_backends.ipynb
@@ -27,7 +27,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 46,
+   "execution_count": 1,
    "id": "708f6560",
    "metadata": {},
    "outputs": [],
@@ -58,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 47,
+   "execution_count": 2,
    "id": "d6ae5a1c",
    "metadata": {},
    "outputs": [],
@@ -72,7 +72,7 @@
    "metadata": {},
    "source": [
     "### StoredDocument\n",
-    "The documents are represented in storage as [StoredDocument](narrow_down.data_types.StoredDocument)\n",
+    "The documents are represented in storage as [StoredDocument](narrow_down.storage.StoredDocument)\n",
     "objects. This is a simple dataclass with the following attributes:\n",
     "    \n",
     "\n",
@@ -98,7 +98,7 @@
     "### Storage levels\n",
     "Depending on the usecase, different levels of persistence may be preferable. Sometimes it is enough to only store just enough to be able to get the IDs of matching documents. In other cases, it can be better to store the whole documents or some additional data in the index. This way it is not necessary to have a second database for this data.\n",
     "\n",
-    "The available storage levels as defined in the enum [StorageLevel](narrow_down.data_types.StorageLevel) are:\n",
+    "The available storage levels as defined in the enum [StorageLevel](narrow_down.storage.StorageLevel) are:\n",
     "\n",
     "| Storage level | Effect                                                                                                                      |\n",
     "|:--------------|:----------------------------------------------------------------------------------------------------------------------------|\n",
@@ -114,7 +114,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 48,
+   "execution_count": 3,
    "id": "1809bcdc",
    "metadata": {},
    "outputs": [
@@ -124,7 +124,7 @@
        "StoredDocument(id_=1, document=None, exact_part=None, fingerprint=None, data='additional data')"
       ]
      },
-     "execution_count": 48,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -151,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 49,
+   "execution_count": 5,
    "id": "7c9c09f3",
    "metadata": {},
    "outputs": [
@@ -161,13 +161,13 @@
        "StoredDocument(id_=1, document='the document text', exact_part='the exact part', fingerprint=None, data='additional data')"
       ]
      },
-     "execution_count": 49,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "from narrow_down.data_types import StorageLevel\n",
+    "from narrow_down.storage import StorageLevel\n",
     "\n",
     "storage_backend = InMemoryStore()\n",
     "\n",
@@ -211,7 +211,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 6,
    "id": "b2cea474",
    "metadata": {},
    "outputs": [
@@ -221,7 +221,7 @@
        "StoredDocument(id_=1, document='the document text', exact_part='the exact part', fingerprint=None, data='additional data')"
       ]
      },
-     "execution_count": 50,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -249,7 +249,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 7,
    "id": "bc58ec9c",
    "metadata": {
     "tags": [
@@ -284,7 +284,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": 8,
    "id": "c4648d39",
    "metadata": {},
    "outputs": [],
@@ -317,7 +317,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": 9,
    "id": "2f85ba8a",
    "metadata": {},
    "outputs": [],
@@ -349,7 +349,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": 10,
    "id": "baf6f6bf",
    "metadata": {
     "tags": [
@@ -363,7 +363,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 11,
    "id": "637d11ab",
    "metadata": {},
    "outputs": [
@@ -373,7 +373,7 @@
        "StoredDocument(id_=1, document='the document text', exact_part='the exact part', fingerprint=None, data='additional data')"
       ]
      },
-     "execution_count": 55,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -400,7 +400,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 56,
+   "execution_count": 12,
    "id": "599899dd",
    "metadata": {
     "tags": [
diff --git a/narrow_down/__init__.py b/narrow_down/__init__.py
index 31f1631..aa227b1 100644
--- a/narrow_down/__init__.py
+++ b/narrow_down/__init__.py
@@ -4,4 +4,4 @@
 __email__ = "chr1st1ank@krudewig-online.de"
 __version__ = "0.8.0"
 
-from . import data_types, hash, similarity_store, storage  # noqa
+from . import hash, similarity_store, storage  # noqa
diff --git a/narrow_down/_minhash.py b/narrow_down/_minhash.py
index 636b6a6..1792d40 100644
--- a/narrow_down/_minhash.py
+++ b/narrow_down/_minhash.py
@@ -17,8 +17,7 @@
 from scipy.integrate import quad as integrate
 
 from . import _rust, hash
-from .data_types import Fingerprint, StorageLevel, StoredDocument, TooLowStorageLevel
-from .storage import StorageBackend
+from .storage import Fingerprint, StorageBackend, StorageLevel, StoredDocument, TooLowStorageLevel
 
 _MERSENNE_PRIME = np.uint32((1 << 32) - 1)
 
diff --git a/narrow_down/data_types.py b/narrow_down/data_types.py
deleted file mode 100644
index 0263674..0000000
--- a/narrow_down/data_types.py
+++ /dev/null
@@ -1,105 +0,0 @@
-"""Type definitions of pure data classes and abstract types."""
-import dataclasses
-import enum
-import pickle  # noqa
-from dataclasses import dataclass
-from typing import Any, Dict, NewType, Optional
-
-import numpy as np
-from numpy import typing as npt
-
-from narrow_down.proto.stored_document_pb2 import StoredDocumentProto
-
-
-class TooLowStorageLevel(Exception):
-    """Raised if a feature is used for which a higher storage level is needed."""
-
-
-class AlreadyInitialized(Exception):
-    """Raised when initializing storage twice or changing immutable settings."""
-
-
-class StorageLevel(enum.Flag):
-    """Detail level of document persistence."""
-
-    Minimal = enum.auto()
-    """Minimal storage level. Only store the necessary data to perform the search."""
-    Fingerprint = enum.auto()
-    """In addition to Minimal, also store the fingerprint, e.g. the Minhashes"""
-    Document = enum.auto()
-    """Store the whole inserted document internally."""
-    Full = Minimal | Fingerprint | Document
-    """Store everything."""
-
-
-Fingerprint = NewType("Fingerprint", npt.NDArray[np.uint32])
-"""Type representing the result of a minhashing operation"""
-
-
-@dataclass(frozen=True)
-class StoredDocument:
-    """Data object combining all possible fields of a document stored."""
-
-    id_: Optional[int] = None
-    """Identifier used to distinguish the document from an identical one."""
-
-    document: Optional[str] = None
-    """The actual content to use for fuzzy matching, e.g. a full unprocessed sentence."""
-
-    exact_part: Optional[str] = None
-    """A string which should be matched exactly."""
-
-    fingerprint: Optional[Fingerprint] = None
-    """A fuzzy fingerprint of the document, e.g. a Minhash."""
-
-    data: Optional[str] = None
-    """Payload to persist together with the document in the internal data structures."""
-
-    def serialize(self, storage_level: StorageLevel) -> bytes:
-        """Serialize a document to bytes."""
-        return StoredDocumentProto(
-            fingerprint=list(self.fingerprint)
-            if self.fingerprint is not None and storage_level & StorageLevel.Fingerprint
-            else [],
-            **{f: getattr(self, f) for f in _FIELDS_FOR_STORAGE_LEVEL[storage_level]},
-        ).SerializeToString()
-
-    @staticmethod
-    def deserialize(doc: bytes, id_: int) -> "StoredDocument":
-        """Deserialize a document from bytes."""
-        p = StoredDocumentProto.FromString(doc)  # type: ignore
-        args: Dict[str, Any] = dict(
-            id_=id_,
-        )
-        if p.HasField("document"):
-            args["document"] = p.document
-        if p.HasField("exact_part"):
-            args["exact_part"] = p.exact_part
-        if p.fingerprint:
-            args["fingerprint"] = Fingerprint(np.array(p.fingerprint, dtype=np.uint32))
-        if p.HasField("data"):
-            args["data"] = p.data
-        return StoredDocument(**args)
-
-    def without(self, *attributes: str) -> "StoredDocument":
-        """Create a copy with the specified attributes left out.
-
-        Args:
-            attributes: The names of the attributes to leave empty
-
-        Returns:
-            A copy of the StoredDocument with all the attributes specified in attributes left out.
-            So they will have their default value (None).
-        """
-        return StoredDocument(
-            **{k: v for k, v in dataclasses.asdict(self).items() if k not in attributes}
-        )
-
-
-_FIELDS_FOR_STORAGE_LEVEL = {
-    StorageLevel.Minimal: {"data"},
-    StorageLevel.Document: {"data", "document", "exact_part"},
-    StorageLevel.Fingerprint: {"data", "exact_part"},
-    StorageLevel.Full: {"data", "document", "exact_part"},
-}
-"""Fields of StoredDocument which need to be serialized to reach a certain storage level."""
diff --git a/narrow_down/scylladb.py b/narrow_down/scylladb.py
index 64615f5..6079da5 100644
--- a/narrow_down/scylladb.py
+++ b/narrow_down/scylladb.py
@@ -95,9 +95,6 @@ async def initialize(
 
         Returns:
             self
-
-        Raises:
-            AlreadyInitialized: If the database file is already initialized
         """
         # Note: CQL does not know unsigned integers.
         #  So we need to take the 64bit (signed) bigint to hold a 32bit unsigned int safely.
diff --git a/narrow_down/similarity_store.py b/narrow_down/similarity_store.py
index f97e256..9e66fdb 100644
--- a/narrow_down/similarity_store.py
+++ b/narrow_down/similarity_store.py
@@ -5,8 +5,13 @@
 
 from narrow_down import _minhash, _tokenize
 from narrow_down._minhash import MinhashLshConfig
-from narrow_down.data_types import StorageLevel, StoredDocument, TooLowStorageLevel
-from narrow_down.storage import InMemoryStore, StorageBackend
+from narrow_down.storage import (
+    InMemoryStore,
+    StorageBackend,
+    StorageLevel,
+    StoredDocument,
+    TooLowStorageLevel,
+)
 
 
 class SimilarityStore:
@@ -281,7 +286,7 @@ async def query(
                 validation is done if the data is available, otherwise not.
 
         Returns:
-            A List of :obj:`~narrow_down.data_types.StoredDocument` objects with all elements
+            A List of :obj:`~narrow_down.storage.StoredDocument` objects with all elements
             which are estimated to be above the similarity threshold.
         """
         tokens = self._tokenize_callable(document)
@@ -305,7 +310,7 @@ async def query_top_n(
                 validation is done if the data is available, otherwise not.
 
         Returns:
-            A List of :obj:`~narrow_down.data_types.StoredDocument` objects with the n
+            A List of :obj:`~narrow_down.storage.StoredDocument` objects with the n
             elements which are most likely above the similarity threshold.
 
         Note that the results are probabilistic. The documents are assumed to be the most likely
diff --git a/narrow_down/sqlite.py b/narrow_down/sqlite.py
index 2c025fc..75716d6 100644
--- a/narrow_down/sqlite.py
+++ b/narrow_down/sqlite.py
@@ -2,7 +2,6 @@
 import sqlite3
 from typing import Iterable, Optional
 
-from narrow_down.data_types import AlreadyInitialized
 from narrow_down.storage import StorageBackend
 
 
@@ -24,27 +23,25 @@ async def initialize(
 
         Returns:
             self
-
-        Raises:
-            AlreadyInitialized: If the database file is already initialized
         """
-        try:
-            with self._connection as conn:
-                conn.execute("CREATE TABLE settings (key TEXT NOT NULL PRIMARY KEY, value TEXT)")
-                conn.execute("CREATE TABLE documents (id INTEGER NOT NULL PRIMARY KEY, doc BLOB)")
-                for i in range(self.partitions):
-                    conn.execute(
-                        f"CREATE TABLE buckets_{i} ("
-                        "bucket INTEGER NOT NULL, "
-                        "hash INTEGER NOT NULL, "
-                        "doc_id INTEGER NOT NULL"
-                        ")"
-                    )
-                conn.execute("PRAGMA synchronous = OFF")
-                conn.execute("PRAGMA journal_mode = MEMORY")
-                conn.commit()
-        except sqlite3.OperationalError as e:
-            raise AlreadyInitialized from e
+        with self._connection as conn:
+            conn.execute(
+                "CREATE TABLE IF NOT EXISTS settings (key TEXT NOT NULL PRIMARY KEY, value TEXT)"
+            )
+            conn.execute(
+                "CREATE TABLE IF NOT EXISTS documents (id INTEGER NOT NULL PRIMARY KEY, doc BLOB)"
+            )
+            for i in range(self.partitions):
+                conn.execute(
+                    f"CREATE TABLE IF NOT EXISTS buckets_{i} ("
+                    "bucket INTEGER NOT NULL, "
+                    "hash INTEGER NOT NULL, "
+                    "doc_id INTEGER NOT NULL"
+                    ")"
+                )
+            conn.execute("PRAGMA synchronous = OFF")
+            conn.execute("PRAGMA journal_mode = MEMORY")
+            conn.commit()
 
         await self.insert_setting("__sqlite_partitions", str(self.partitions))
 
diff --git a/narrow_down/storage.py b/narrow_down/storage.py
index 41068a9..cc0d9eb 100644
--- a/narrow_down/storage.py
+++ b/narrow_down/storage.py
@@ -1,10 +1,109 @@
 """Base classes and interfaces for storage."""
+import dataclasses
+import enum
+import pickle  # noqa
 from abc import ABC
-from typing import Iterable, Optional
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, NewType, Optional
+
+import numpy as np
+from numpy import typing as npt
+
+from narrow_down.proto.stored_document_pb2 import StoredDocumentProto
 
 from ._rust import RustMemoryStore
 
 
+class TooLowStorageLevel(Exception):
+    """Raised if a feature is used for which a higher storage level is needed."""
+
+
+class StorageLevel(enum.Flag):
+    """Detail level of document persistence."""
+
+    Minimal = enum.auto()
+    """Minimal storage level. Only store the necessary data to perform the search."""
+    Fingerprint = enum.auto()
+    """In addition to Minimal, also store the fingerprint, e.g. the Minhashes"""
+    Document = enum.auto()
+    """Store the whole inserted document internally."""
+    Full = Minimal | Fingerprint | Document
+    """Store everything."""
+
+
+Fingerprint = NewType("Fingerprint", npt.NDArray[np.uint32])
+"""Type representing the result of a minhashing operation"""
+
+
+@dataclass(frozen=True)
+class StoredDocument:
+    """Data object combining all possible fields of a document stored."""
+
+    id_: Optional[int] = None
+    """Identifier used to distinguish the document from an identical one."""
+
+    document: Optional[str] = None
+    """The actual content to use for fuzzy matching, e.g. a full unprocessed sentence."""
+
+    exact_part: Optional[str] = None
+    """A string which should be matched exactly."""
+
+    fingerprint: Optional[Fingerprint] = None
+    """A fuzzy fingerprint of the document, e.g. a Minhash."""
+
+    data: Optional[str] = None
+    """Payload to persist together with the document in the internal data structures."""
+
+    def serialize(self, storage_level: StorageLevel) -> bytes:
+        """Serialize a document to bytes."""
+        return StoredDocumentProto(
+            fingerprint=list(self.fingerprint)
+            if self.fingerprint is not None and storage_level & StorageLevel.Fingerprint
+            else [],
+            **{f: getattr(self, f) for f in _FIELDS_FOR_STORAGE_LEVEL[storage_level]},
+        ).SerializeToString()
+
+    @staticmethod
+    def deserialize(doc: bytes, id_: int) -> "StoredDocument":
+        """Deserialize a document from bytes."""
+        p = StoredDocumentProto.FromString(doc)  # type: ignore
+        args: Dict[str, Any] = dict(
+            id_=id_,
+        )
+        if p.HasField("document"):
+            args["document"] = p.document
+        if p.HasField("exact_part"):
+            args["exact_part"] = p.exact_part
+        if p.fingerprint:
+            args["fingerprint"] = Fingerprint(np.array(p.fingerprint, dtype=np.uint32))
+        if p.HasField("data"):
+            args["data"] = p.data
+        return StoredDocument(**args)
+
+    def without(self, *attributes: str) -> "StoredDocument":
+        """Create a copy with the specified attributes left out.
+
+        Args:
+            attributes: The names of the attributes to leave empty
+
+        Returns:
+            A copy of the StoredDocument with all the attributes specified in attributes left out.
+            So they will have their default value (None).
+        """
+        return StoredDocument(
+            **{k: v for k, v in dataclasses.asdict(self).items() if k not in attributes}
+        )
+
+
+_FIELDS_FOR_STORAGE_LEVEL = {
+    StorageLevel.Minimal: {"data"},
+    StorageLevel.Document: {"data", "document", "exact_part"},
+    StorageLevel.Fingerprint: {"data", "exact_part"},
+    StorageLevel.Full: {"data", "document", "exact_part"},
+}
+"""Fields of StoredDocument which need to be serialized to reach a certain storage level."""
+
+
 class StorageBackend(ABC):
     """Storage backend for a SimilarityStore."""
 
@@ -15,8 +114,6 @@ async def initialize(
 
         Returns:
             self
-
-        Implementations may raise an error if the backend is already initialized.
         """
         return self
 


Test patch:
diff --git a/tests/test_backend_benchmarks.py b/tests/test_backend_benchmarks.py
index 3aba581..347a666 100644
--- a/tests/test_backend_benchmarks.py
+++ b/tests/test_backend_benchmarks.py
@@ -17,11 +17,10 @@
 import cassandra.cluster  # type: ignore
 import pytest
 
-from narrow_down.data_types import StorageLevel
 from narrow_down.scylladb import ScyllaDBStore
 from narrow_down.similarity_store import SimilarityStore
 from narrow_down.sqlite import SQLiteStore
-from narrow_down.storage import InMemoryStore
+from narrow_down.storage import InMemoryStore, StorageLevel
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_data_types.py b/tests/test_data_types.py
deleted file mode 100644
index 960391c..0000000
--- a/tests/test_data_types.py
+++ /dev/null
@@ -1,49 +0,0 @@
-"""Tests for `narrow_down.data_types`."""
-import dataclasses
-
-import numpy as np
-import pytest
-
-from narrow_down.data_types import Fingerprint, StorageLevel, StoredDocument
-
-
-@pytest.mark.parametrize("data", [None, "", "user data"])
-@pytest.mark.parametrize("fingerprint", [None, Fingerprint(np.array([1]))])
-@pytest.mark.parametrize("exact_part", [None, "", "text_exact_part"])
-@pytest.mark.parametrize("document", [None, "", "text_document"])
-@pytest.mark.parametrize("id_", [None, 0, 1])
-@pytest.mark.parametrize(
-    "storage_level, expected_fields",
-    [
-        (StorageLevel.Minimal, {"data"}),
-        (StorageLevel.Document, {"data", "document", "exact_part"}),
-        (StorageLevel.Fingerprint, {"data", "exact_part", "fingerprint"}),
-        (StorageLevel.Full, {"data", "document", "exact_part", "fingerprint"}),
-    ],
-)
-def test_stored_document_serialization(
-    storage_level, expected_fields, id_, document, exact_part, fingerprint, data
-):
-    document = StoredDocument(
-        id_=id_,
-        document=document,
-        exact_part=exact_part,
-        fingerprint=fingerprint,
-        data=data,
-    )
-    print("before:", document)
-    deserialized = StoredDocument.deserialize(document.serialize(storage_level), id_=None)
-    print("after:", deserialized)
-    for field in dataclasses.fields(deserialized):
-        if field.name in expected_fields:
-            assert getattr(deserialized, field.name) == getattr(document, field.name)
-        else:
-            assert getattr(deserialized, field.name) is None
-
-
-def test_stored_document_without():
-    document = StoredDocument(id_=5, document="abcd")
-    assert document.without("document") == StoredDocument(id_=5)
-    assert document.without("document", "id_") == StoredDocument()
-    assert document.without() == document
-    assert document.without() is not document
diff --git a/tests/test_minhash.py b/tests/test_minhash.py
index 94a8563..9b9ef5c 100644
--- a/tests/test_minhash.py
+++ b/tests/test_minhash.py
@@ -2,8 +2,8 @@
 import numpy as np
 import pytest
 
-from narrow_down import _minhash, data_types, storage
-from narrow_down.data_types import StorageLevel, TooLowStorageLevel
+from narrow_down import _minhash, storage
+from narrow_down.storage import StorageLevel, StoredDocument, TooLowStorageLevel
 
 
 def test_minhash_lsh_config__json():
@@ -37,11 +37,11 @@ def f():
 @pytest.mark.asyncio
 async def test_lsh__basic_lookup_without_exact_part():
     """Minimal check if an LSH can be constructed."""
-    test_doc = data_types.StoredDocument(
+    test_doc = StoredDocument(
         # document="abc def",
         exact_part="exact:part",
         # data="some custom payload",
-        fingerprint=data_types.Fingerprint(np.array([2, 4, 6])),
+        fingerprint=storage.Fingerprint(np.array([2, 4, 6])),
     )
     lsh = _minhash.LSH(
         _minhash.MinhashLshConfig(n_hashes=2, n_bands=2, rows_per_band=1),
@@ -57,11 +57,11 @@ async def test_lsh__basic_lookup_without_exact_part():
 @pytest.mark.asyncio
 async def test_lsh__basic_lookup_with_exact_part():
     """Minimal check if an LSH can be constructed."""
-    test_doc = data_types.StoredDocument(
+    test_doc = StoredDocument(
         # document="abc def",
         exact_part="exact:part",
         # data="some custom payload",
-        fingerprint=data_types.Fingerprint(np.array([2, 4, 6])),
+        fingerprint=storage.Fingerprint(np.array([2, 4, 6])),
     )
     lsh = _minhash.LSH(
         _minhash.MinhashLshConfig(n_hashes=2, n_bands=2, rows_per_band=1),
@@ -83,39 +83,31 @@ async def test_lsh__top_n():
         storage=await storage.InMemoryStore().initialize(),
     )
     await lsh.insert(
-        data_types.StoredDocument(
-            document="1", fingerprint=data_types.Fingerprint(np.array([2, 4, 6]))
-        ),
+        StoredDocument(document="1", fingerprint=storage.Fingerprint(np.array([2, 4, 6]))),
     )
     await lsh.insert(
-        data_types.StoredDocument(
-            document="2", fingerprint=data_types.Fingerprint(np.array([2, 4, 7]))
-        ),
+        StoredDocument(document="2", fingerprint=storage.Fingerprint(np.array([2, 4, 7]))),
     )
     await lsh.insert(
-        data_types.StoredDocument(
-            document="3", fingerprint=data_types.Fingerprint(np.array([2, 5, 7]))
-        ),
+        StoredDocument(document="3", fingerprint=storage.Fingerprint(np.array([2, 5, 7]))),
     )
     await lsh.insert(
-        data_types.StoredDocument(
-            document="4", fingerprint=data_types.Fingerprint(np.array([3, 5, 7]))
-        ),
+        StoredDocument(document="4", fingerprint=storage.Fingerprint(np.array([3, 5, 7]))),
     )
 
-    result = await lsh.query_top_n(n=1, fingerprint=data_types.Fingerprint(np.array([2, 4, 6])))
+    result = await lsh.query_top_n(n=1, fingerprint=storage.Fingerprint(np.array([2, 4, 6])))
     print(result)
     assert sorted([r.document for r in result]) == ["1"]
 
-    result = await lsh.query_top_n(n=2, fingerprint=data_types.Fingerprint(np.array([2, 4, 6])))
+    result = await lsh.query_top_n(n=2, fingerprint=storage.Fingerprint(np.array([2, 4, 6])))
     print(result)
     assert sorted([r.document for r in result]) == ["1", "2"]
 
-    result = await lsh.query_top_n(n=3, fingerprint=data_types.Fingerprint(np.array([2, 4, 6])))
+    result = await lsh.query_top_n(n=3, fingerprint=storage.Fingerprint(np.array([2, 4, 6])))
     print(result)
     assert sorted([r.document for r in result]) == ["1", "2", "3"]
 
-    result = await lsh.query_top_n(n=4, fingerprint=data_types.Fingerprint(np.array([2, 4, 6])))
+    result = await lsh.query_top_n(n=4, fingerprint=storage.Fingerprint(np.array([2, 4, 6])))
     print(result)
     assert sorted([r.document for r in result]) == ["1", "2", "3"]
 
@@ -125,14 +117,14 @@ async def test_lsh__insert_invalid_document():
     """Test error handling of insert()."""
     lsh = _minhash.LSH(_minhash.MinhashLshConfig(1, 1, 1), None)
     with pytest.raises(ValueError):
-        await lsh.insert(data_types.StoredDocument())
+        await lsh.insert(StoredDocument())
 
 
 @pytest.mark.asyncio
 async def test_lsh__remove_by_id():
     """Try to remove a document from the index."""
-    f = data_types.Fingerprint(np.array([2, 4, 6]))
-    test_doc = data_types.StoredDocument(fingerprint=f)
+    f = storage.Fingerprint(np.array([2, 4, 6]))
+    test_doc = StoredDocument(fingerprint=f)
     lsh = _minhash.LSH(
         _minhash.MinhashLshConfig(n_hashes=2, n_bands=2, rows_per_band=1),
         storage=await storage.InMemoryStore().initialize(),
@@ -145,8 +137,8 @@ async def test_lsh__remove_by_id():
 @pytest.mark.asyncio
 async def test_lsh__remove_by_id__missing_fingerprint():
     """Try to remove a document from the index when the storage level is too low."""
-    f = data_types.Fingerprint(np.array([2, 4, 6]))
-    test_doc = data_types.StoredDocument(fingerprint=f)
+    f = storage.Fingerprint(np.array([2, 4, 6]))
+    test_doc = StoredDocument(fingerprint=f)
     lsh = _minhash.LSH(
         _minhash.MinhashLshConfig(n_hashes=2, n_bands=2, rows_per_band=1),
         storage=await storage.InMemoryStore().initialize(),
diff --git a/tests/test_similarity_store.py b/tests/test_similarity_store.py
index 6165956..54ceb99 100644
--- a/tests/test_similarity_store.py
+++ b/tests/test_similarity_store.py
@@ -1,10 +1,10 @@
 """Tests for `narrow_down.similarity_store`."""
 import pytest
 
-import narrow_down.data_types
-from narrow_down.data_types import StorageLevel, StoredDocument
+import narrow_down.storage
 from narrow_down.similarity_store import SimilarityStore
 from narrow_down.sqlite import SQLiteStore
+from narrow_down.storage import StorageLevel, StoredDocument
 
 
 @pytest.mark.asyncio
@@ -303,7 +303,7 @@ async def test_similarity_store__remove_by_id__error_storage_level():
     sample_doc = "Some example document"
 
     doc_id = await simstore.insert(sample_doc)
-    with pytest.raises(narrow_down.data_types.TooLowStorageLevel):
+    with pytest.raises(narrow_down.storage.TooLowStorageLevel):
         await simstore.remove_by_id(doc_id)
     results = await simstore.query(sample_doc)
 
diff --git a/tests/test_sqlite.py b/tests/test_sqlite.py
index d1f50ed..b9dd1f6 100644
--- a/tests/test_sqlite.py
+++ b/tests/test_sqlite.py
@@ -52,8 +52,9 @@ async def insert_setting():
 async def test_sqlite_store__init__already_exists(tmp_path):
     dbfile = str(tmp_path / "test.db")
     await narrow_down.sqlite.SQLiteStore(dbfile).initialize()
-    with pytest.raises(narrow_down.data_types.AlreadyInitialized):
-        await narrow_down.sqlite.SQLiteStore(dbfile).initialize()
+    ims = await narrow_down.sqlite.SQLiteStore(dbfile).initialize()
+    await ims.insert_setting(key="k", value="155")
+    assert await ims.query_setting("k") == "155"
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 82ac911..777698f 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -1,7 +1,52 @@
 """Tests for the `narrow_down.storage` module."""
+import dataclasses
+
+import numpy as np
 import pytest
 
-from narrow_down.storage import InMemoryStore
+from narrow_down.storage import Fingerprint, InMemoryStore, StorageLevel, StoredDocument
+
+
+@pytest.mark.parametrize("data", [None, "", "user data"])
+@pytest.mark.parametrize("fingerprint", [None, Fingerprint(np.array([1]))])
+@pytest.mark.parametrize("exact_part", [None, "", "text_exact_part"])
+@pytest.mark.parametrize("document", [None, "", "text_document"])
+@pytest.mark.parametrize("id_", [None, 0, 1])
+@pytest.mark.parametrize(
+    "storage_level, expected_fields",
+    [
+        (StorageLevel.Minimal, {"data"}),
+        (StorageLevel.Document, {"data", "document", "exact_part"}),
+        (StorageLevel.Fingerprint, {"data", "exact_part", "fingerprint"}),
+        (StorageLevel.Full, {"data", "document", "exact_part", "fingerprint"}),
+    ],
+)
+def test_stored_document_serialization(
+    storage_level, expected_fields, id_, document, exact_part, fingerprint, data
+):
+    document = StoredDocument(
+        id_=id_,
+        document=document,
+        exact_part=exact_part,
+        fingerprint=fingerprint,
+        data=data,
+    )
+    print("before:", document)
+    deserialized = StoredDocument.deserialize(document.serialize(storage_level), id_=None)
+    print("after:", deserialized)
+    for field in dataclasses.fields(deserialized):
+        if field.name in expected_fields:
+            assert getattr(deserialized, field.name) == getattr(document, field.name)
+        else:
+            assert getattr(deserialized, field.name) is None
+
+
+def test_stored_document_without():
+    document = StoredDocument(id_=5, document="abcd")
+    assert document.without("document") == StoredDocument(id_=5)
+    assert document.without("document", "id_") == StoredDocument()
+    assert document.without() == document
+    assert document.without() is not document
 
 
 @pytest.mark.asyncio