feat: add NumPy array serialization support in SQLSpec plugin (#165)

Adds automatic NumPy array serialization to SQLSpec's Litestar plugin, enabling seamless bidirectional conversion between NumPy arrays and JSON for vector embedding workflows.

Author: cofin

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

GEMINI.md

@@ -25,8 +25,9 @@
     pool_provider_maker,
     session_provider_maker,
 )
-from sqlspec.typing import ConnectionT, PoolT
+from sqlspec.typing import NUMPY_INSTALLED, ConnectionT, PoolT, SchemaT
 from sqlspec.utils.logging import get_logger
+from sqlspec.utils.serializers import numpy_array_dec_hook, numpy_array_enc_hook, numpy_array_predicate
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable
@@ -82,6 +83,10 @@ class _PluginConfigState:
 class SQLSpecPlugin(InitPluginProtocol, CLIPlugin):
     """Litestar plugin for SQLSpec database integration.
 
+    Automatically configures NumPy array serialization when NumPy is installed,
+    enabling seamless bidirectional conversion between NumPy arrays and JSON
+    for vector embedding workflows.
+
     Session Table Migrations:
         The Litestar extension includes migrations for creating session storage tables.
         To include these migrations in your database migration workflow, add 'litestar'
@@ -225,6 +230,8 @@ def on_cli_init(self, cli: "Group") -> None:
     def on_app_init(self, app_config: "AppConfig") -> "AppConfig":
         """Configure Litestar application with SQLSpec database integration.
 
+        Automatically registers NumPy array serialization when NumPy is installed.
+
         Args:
             app_config: The Litestar application configuration instance.
 
@@ -239,7 +246,7 @@ def store_sqlspec_in_state() -> None:
         app_config.on_startup.append(store_sqlspec_in_state)
         app_config.signature_types.extend([SQLSpec, DatabaseConfigProtocol, SyncConfigT, AsyncConfigT])
 
-        signature_namespace = {"ConnectionT": ConnectionT, "PoolT": PoolT, "DriverT": DriverT}
+        signature_namespace = {"ConnectionT": ConnectionT, "PoolT": PoolT, "DriverT": DriverT, "SchemaT": SchemaT}
 
         for state in self._plugin_configs:
             state.annotation = type(state.config)
@@ -262,6 +269,23 @@ def store_sqlspec_in_state() -> None:
         if signature_namespace:
             app_config.signature_namespace.update(signature_namespace)
 
+        if NUMPY_INSTALLED:
+            import numpy as np
+
+            if app_config.type_encoders is None:
+                app_config.type_encoders = {np.ndarray: numpy_array_enc_hook}
+            else:
+                encoders_dict = dict(app_config.type_encoders)
+                encoders_dict[np.ndarray] = numpy_array_enc_hook
+                app_config.type_encoders = encoders_dict
+
+            if app_config.type_decoders is None:
+                app_config.type_decoders = [(numpy_array_predicate, numpy_array_dec_hook)]  # type: ignore[list-item]
+            else:
+                decoders_list = list(app_config.type_decoders)
+                decoders_list.append((numpy_array_predicate, numpy_array_dec_hook))  # type: ignore[arg-type]
+                app_config.type_decoders = decoders_list
+
         return app_config
 
     def get_annotations(

GEMINI.md

@@ -2,11 +2,15 @@
 
 Re-exports common JSON encoding and decoding functions from the core
 serialization module for convenient access.
+
+Provides NumPy array serialization hooks for framework integrations
+that support custom type encoders and decoders (e.g., Litestar).
 """
 
 from typing import Any, Literal, overload
 
 from sqlspec._serialization import decode_json, encode_json
+from sqlspec.typing import NUMPY_INSTALLED
 
 
 @overload
@@ -55,4 +59,109 @@ def from_json(data: str | bytes, *, decode_bytes: bool = True) -> Any:
     return decode_json(data)
 
 
-__all__ = ("from_json", "to_json")
+def numpy_array_enc_hook(value: Any) -> Any:
+    """Encode NumPy array to JSON-compatible list.
+
+    Converts NumPy ndarrays to Python lists for JSON serialization.
+    Gracefully handles cases where NumPy is not installed by returning
+    the original value unchanged.
+
+    Args:
+        value: Value to encode (checked for ndarray type).
+
+    Returns:
+        List representation if value is ndarray, original value otherwise.
+
+    Example:
+        >>> import numpy as np
+        >>> arr = np.array([1.0, 2.0, 3.0])
+        >>> numpy_array_enc_hook(arr)
+        [1.0, 2.0, 3.0]
+
+        >>> # Multi-dimensional arrays work automatically
+        >>> arr_2d = np.array([[1, 2], [3, 4]])
+        >>> numpy_array_enc_hook(arr_2d)
+        [[1, 2], [3, 4]]
+    """
+    if not NUMPY_INSTALLED:
+        return value
+
+    import numpy as np
+
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return value
+
+
+def numpy_array_dec_hook(value: Any) -> "Any":
+    """Decode list to NumPy array.
+
+    Converts Python lists to NumPy arrays when appropriate.
+    Works best with typed schemas (Pydantic, msgspec) that expect ndarray.
+
+    Args:
+        value: List to potentially convert to ndarray.
+
+    Returns:
+        NumPy array if conversion successful, original value otherwise.
+
+    Note:
+        Dtype is inferred by NumPy and may differ from original array.
+        For explicit dtype control, construct arrays manually in application code.
+
+    Example:
+        >>> numpy_array_dec_hook([1.0, 2.0, 3.0])
+        array([1., 2., 3.])
+
+        >>> # Returns original value if NumPy not installed
+        >>> # (when NUMPY_INSTALLED is False)
+        >>> numpy_array_dec_hook([1, 2, 3])
+        [1, 2, 3]
+    """
+    if not NUMPY_INSTALLED:
+        return value
+
+    import numpy as np
+
+    if isinstance(value, list):
+        try:
+            return np.array(value)
+        except Exception:
+            return value
+    return value
+
+
+def numpy_array_predicate(value: Any) -> bool:
+    """Check if value is NumPy array instance.
+
+    Type checker for decoder registration in framework plugins.
+    Returns False when NumPy is not installed.
+
+    Args:
+        value: Value to type-check.
+
+    Returns:
+        True if value is ndarray, False otherwise.
+
+    Example:
+        >>> import numpy as np
+        >>> numpy_array_predicate(np.array([1, 2, 3]))
+        True
+
+        >>> numpy_array_predicate([1, 2, 3])
+        False
+
+        >>> # Returns False when NumPy not installed
+        >>> # (when NUMPY_INSTALLED is False)
+        >>> numpy_array_predicate([1, 2, 3])
+        False
+    """
+    if not NUMPY_INSTALLED:
+        return False
+
+    import numpy as np
+
+    return isinstance(value, np.ndarray)
+
+
+__all__ = ("from_json", "numpy_array_dec_hook", "numpy_array_enc_hook", "numpy_array_predicate", "to_json")