feat: add optional numpy serialization support (#104)

Implements optional NumPy array serialization support in SQLSpec's serialization system

Author: cofin

sqlspec/_serialization.py

@@ -2,6 +2,9 @@
 
 Provides a Protocol-based serialization system that users can extend.
 Supports msgspec, orjson, and standard library JSON with automatic fallback.
+
+Features optional numpy array serialization when numpy is installed.
+Arrays are automatically converted to lists during JSON encoding.
 """
 
 import contextlib
@@ -11,17 +14,23 @@
 from abc import ABC, abstractmethod
 from typing import Any, Final, Literal, Protocol, overload
 
+from sqlspec._typing import NUMPY_INSTALLED
 from sqlspec.typing import MSGSPEC_INSTALLED, ORJSON_INSTALLED, PYDANTIC_INSTALLED, BaseModel
 
 
-def _type_to_string(value: Any) -> str:  # pragma: no cover
+def _type_to_string(value: Any) -> Any:  # pragma: no cover
     """Convert special types to strings for JSON serialization.
 
+    Handles datetime, date, enums, Pydantic models, and numpy arrays.
+
     Args:
         value: Value to convert.
 
     Returns:
-        String representation of the value.
+        Serializable representation of the value (string, list, dict, etc.).
+
+    Raises:
+        TypeError: If value cannot be serialized.
     """
     if isinstance(value, datetime.datetime):
         return convert_datetime_to_gmt_iso(value)
@@ -31,10 +40,16 @@ def _type_to_string(value: Any) -> str:  # pragma: no cover
         return str(value.value)
     if PYDANTIC_INSTALLED and isinstance(value, BaseModel):
         return value.model_dump_json()
+    if NUMPY_INSTALLED:
+        import numpy as np
+
+        if isinstance(value, np.ndarray):
+            return value.tolist()
     try:
         return str(value)
     except Exception as exc:
-        raise TypeError from exc
+        msg = f"Cannot serialize {type(value).__name__}"
+        raise TypeError(msg) from exc
 
 
 class JSONSerializer(Protocol):
@@ -128,22 +143,37 @@ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
 
 
 class OrjsonSerializer(BaseJSONSerializer):
-    """Orjson-based JSON serializer with native datetime/UUID support."""
+    """Orjson-based JSON serializer with native datetime/UUID support.
+
+    Automatically enables numpy serialization if numpy is installed.
+    """
 
     __slots__ = ()
 
     def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
-        """Encode data using orjson."""
+        """Encode data using orjson.
+
+        Args:
+            data: Data to encode.
+            as_bytes: Whether to return bytes instead of string.
+
+        Returns:
+            JSON string or bytes depending on as_bytes parameter.
+        """
         from orjson import (
             OPT_NAIVE_UTC,  # pyright: ignore[reportUnknownVariableType]
-            OPT_SERIALIZE_NUMPY,  # pyright: ignore[reportUnknownVariableType]
             OPT_SERIALIZE_UUID,  # pyright: ignore[reportUnknownVariableType]
         )
         from orjson import dumps as _orjson_dumps  # pyright: ignore[reportMissingImports]
 
-        result = _orjson_dumps(
-            data, default=_type_to_string, option=OPT_SERIALIZE_NUMPY | OPT_NAIVE_UTC | OPT_SERIALIZE_UUID
-        )
+        options = OPT_NAIVE_UTC | OPT_SERIALIZE_UUID
+
+        if NUMPY_INSTALLED:
+            from orjson import OPT_SERIALIZE_NUMPY  # pyright: ignore[reportUnknownVariableType]
+
+            options |= OPT_SERIALIZE_NUMPY
+
+        result = _orjson_dumps(data, default=_type_to_string, option=options)
         return result if as_bytes else result.decode("utf-8")
 
     def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:

sqlspec/_typing.py

@@ -389,6 +389,24 @@ def slice(self, offset: int = 0, length: "int | None" = None) -> Any:
     PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
 
 
+@runtime_checkable
+class NumpyArrayStub(Protocol):
+    """Protocol stub for numpy.ndarray when numpy is not installed.
+
+    Provides minimal interface for type checking and serialization support.
+    """
+
+    def tolist(self) -> "list[Any]":
+        """Convert array to Python list."""
+        ...
+
+
+try:
+    from numpy import ndarray as NumpyArray  # noqa: N812
+except ImportError:
+    NumpyArray = NumpyArrayStub  # type: ignore[assignment,misc]
+
+
 try:
     from opentelemetry import trace  # pyright: ignore[reportMissingImports, reportAssignmentType]
     from opentelemetry.trace import (  # pyright: ignore[reportMissingImports, reportAssignmentType]
@@ -660,6 +678,8 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter
     "FailFastStub",
     "Gauge",
     "Histogram",
+    "NumpyArray",
+    "NumpyArrayStub",
     "Span",
     "Status",
     "StatusCode",

sqlspec/typing.py

@@ -41,6 +41,7 @@
     FailFast,
     Gauge,
     Histogram,
+    NumpyArray,
     Span,
     Status,
     StatusCode,
@@ -233,6 +234,7 @@ class StorageMixin(MixinOf(DriverProtocol)): ...
     "ModelDictList",
     "ModelDictList",
     "ModelT",
+    "NumpyArray",
     "PoolT",
     "PoolT_co",
     "PydanticOrMsgspecT",