chore: import upstream snapshot with attribution

python/zvec/model/doc.py

@@ -0,0 +1,178 @@
+# Copyright 2025-present the zvec project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import json
+from typing import Any, Optional
+
+from ..common import VectorType
+
+__all__ = [
+    "Doc",
+    "DocList",
+]
+
+
+class Doc:
+    """Represents a retrieved document with optional metadata, fields, and vectors.
+
+    This immutable data class encapsulates the result of a search or retrieval
+    operation. It includes the document ID, relevance score (if applicable),
+    scalar fields, and vector embeddings.
+
+    During initialization, any `numpy.ndarray` in `vectors` is automatically
+    converted to a plain Python list for JSON serialization and immutability.
+
+    Attributes:
+        id (str): Unique identifier of the document.
+        score (Optional[float], optional): Relevance score from search.
+            Defaults to None.
+        vectors (Optional[dict[str, VectorType]], optional): Named vector
+            embeddings associated with the document. Values are converted to
+            lists if originally `np.ndarray`. Defaults to None.
+        fields (Optional[dict[str, Any]], optional): Scalar metadata fields
+            (e.g., title, timestamp). Defaults to None.
+
+    Examples:
+        >>> import numpy as np
+        >>> import zvec
+        >>> doc = zvec.Doc(
+        ...     id="doc1",
+        ...     score=0.95,
+        ...     vectors={"emb": np.array([0.1, 0.2, 0.3])},
+        ...     fields={"title": "Hello World"}
+        ... )
+        >>> print(doc.vector("emb"))
+        [0.1, 0.2, 0.3]
+        >>> print(doc.has_field("title"))
+        True
+    """
+
+    __slots__ = ("id", "score", "vectors", "fields")
+
+    def __init__(
+        self,
+        id: str,
+        score: Optional[float] = None,
+        vectors: Optional[dict[str, VectorType]] = None,
+        fields: Optional[dict[str, Any]] = None,
+    ):
+        self.id = id
+        self.score = score
+        self.vectors = vectors or {}
+        self.fields = fields or {}
+
+    def has_field(self, name: str) -> bool:
+        """Check if the document contains a scalar field with the given name.
+
+        Args:
+            name (str): Name of the field to check.
+
+        Returns:
+            bool: True if the field exists, False otherwise.
+        """
+        return name in self.fields
+
+    def has_vector(self, name: str) -> bool:
+        """Check if the document contains a vector with the given name.
+
+        Args:
+            name (str): Name of the vector to check.
+
+        Returns:
+            bool: True if the vector exists, False otherwise.
+        """
+        return name in self.vectors
+
+    def vector(self, name: str):
+        """Get a vector by name.
+
+        Args:
+            name (str): Name of the vector.
+
+        Returns:
+            Any: The vector (as a list) if it exists, otherwise None.
+        """
+        return self.vectors and self.vectors.get(name)
+
+    def field(self, name: str):
+        """Get a scalar field by name.
+
+        Args:
+            name (str): Name of the field.
+
+        Returns:
+            Any: The field value if it exists, otherwise None.
+        """
+        return self.fields and self.fields.get(name)
+
+    def vector_names(self) -> list[str]:
+        """Get the list of all vector names in this document.
+
+        Returns:
+            list[str]: A list of vector field names. Empty if no vectors.
+        """
+        return [] if not self.vectors else list(self.vectors.keys())
+
+    def field_names(self) -> list[str]:
+        """Get the list of all scalar field names in this document.
+
+        Returns:
+            list[str]: A list of field names. Empty if no fields.
+        """
+        return [] if not self.fields else list(self.fields.keys())
+
+    def __repr__(self) -> str:
+        try:
+            schema = {
+                "id": self.id,
+                "score": self.score,
+                "fields": self.fields,
+                "vectors": self.vectors,
+            }
+            return json.dumps(schema, indent=2, ensure_ascii=False)
+        except Exception as e:
+            return f"<Doc error during repr: {e}>"
+
+    def _replace(self, **changes):
+        new_tuple = (
+            changes.get("id", self.id),
+            changes.get("score", self.score),
+            changes.get("fields", self.fields.copy() if self.fields else None),
+            changes.get("vectors", self.vectors.copy() if self.vectors else None),
+        )
+        return type(self)._from_tuple(new_tuple)
+
+    @classmethod
+    def _from_tuple(
+        cls, data_tuple: tuple[str, float, dict[str, Any], dict[str, VectorType]]
+    ):
+        obj = object.__new__(cls)
+        obj.id = data_tuple[0]
+        obj.score = data_tuple[1]
+        obj.fields = data_tuple[2] or {}
+
+        vectors = data_tuple[3]
+        if vectors is not None:
+            obj.vectors = {
+                name: (vec.tolist() if hasattr(vec, "tolist") else vec)
+                for name, vec in vectors.items()
+            }
+        else:
+            obj.vectors = {}
+        return obj
+
+
+#: Type alias for query results: a list of documents returned by a single query route.
+DocList = list[Doc]