feat(spm): initial REST API client for CryoSPARC

.pre-commit-config.yaml

@@ -27,4 +27,5 @@ repos:
     rev: v1.1.386
     hooks:
       - id: pyright
-        additional_dependencies: [cython, httpretty, numpy, pytest, setuptools]
+        additional_dependencies:
+          [cython, httpretty, httpx, numpy, pydantic, pytest, setuptools]

cryosparc/dataset.py

@@ -146,6 +146,7 @@ class Dataset(Streamable, MutableMapping[str, Column], Generic[R]):
 
     __slots__ = ("_row_class", "_rows", "_data")
 
+    media_type = "application/x-cryosparc-dataset"
     _row_class: Type[R]
     _rows: Optional[Spool[R]]
     _data: Data

cryosparc/errors.py

@@ -2,17 +2,51 @@
 Definitions for various error classes raised by cryosparc-tools functions
 """
 
-from typing import Any, List, TypedDict
+import json
+from typing import TYPE_CHECKING, Any, List, TypedDict
 
 from .spec import Datafield, Datatype, SlotSpec
 
+if TYPE_CHECKING:
+    from httpx import Response
+
 
 class DatasetLoadError(Exception):
     """Exception type raised when a dataset cannot be loaded"""
 
     pass
 
 
+class APIError(ValueError):
+    """
+    Raised by failed request to a CryoSPARC API server.
+    """
+
+    code: int
+    res: "Response"
+    data: Any
+
+    def __init__(
+        self,
+        reason: str,
+        *args: object,
+        res: "Response",
+        data: Any = None,
+    ) -> None:
+        msg = f"*** [API] ({res.request.method} {res.url}, code {res.status_code}) {reason}"
+        super().__init__(msg, *args)
+        self.res = res
+        self.code = res.status_code
+        self.data = data
+
+    def __str__(self):
+        s = super().__str__()
+        if self.data:
+            s += "\n"
+            s += json.dumps(self.data)
+        return s
+
+
 class CommandError(Exception):
     """
     Raised by failed request to a CryoSPARC command server.

cryosparc/json_util.py

@@ -0,0 +1,84 @@
+import base64
+from datetime import datetime
+from pathlib import PurePath
+from typing import Any, Mapping
+
+import numpy as n
+from pydantic import BaseModel
+
+
+def api_default(obj: Any) -> Any:
+    """
+    json.dump "default" argument for sending objects over a JSON API. Ensures
+    that special non-JSON types such as Path are NDArray are encoded correctly.
+    """
+    if isinstance(obj, n.floating):
+        if n.isnan(obj):
+            return float(0)
+        elif n.isposinf(obj):
+            return float("inf")
+        elif n.isneginf(obj):
+            return float("-inf")
+        return float(obj)
+    elif isinstance(obj, n.integer):
+        return int(obj)
+    elif isinstance(obj, n.ndarray):
+        return ndarray_to_json(obj)
+    elif isinstance(obj, bytes):
+        return binary_to_json(obj)
+    elif isinstance(obj, datetime):
+        return obj.isoformat()
+    elif isinstance(obj, PurePath):
+        return str(obj)
+    elif isinstance(obj, BaseModel):
+        return obj.model_dump(exclude_none=True)
+    else:
+        return obj
+
+
+def api_object_hook(dct: Mapping[str, Any]):
+    """
+    json.dump "object_hook" argument for receiving JSON from an API request.
+    Ensures that special objects that are actually encoded numpy arrays or bytes
+    are decoded as such.
+    """
+    if "$ndarray" in dct:
+        return ndarray_from_json(dct)
+    elif "$binary" in dct:
+        return binary_from_json(dct)
+    else:
+        return dct  # pydantic will take care of everything else
+
+
+def binary_to_json(binary: bytes):
+    """
+    Encode bytes as a JSON-serializeable object
+    """
+    return {"$binary": {"base64": base64.b64encode(binary).decode()}}
+
+
+def binary_from_json(dct: Mapping[str, Any]) -> bytes:
+    if "base64" not in dct["$binary"] or not isinstance(b64 := dct["$binary"]["base64"], str):
+        raise TypeError(f"$binary base64 must be a string: {dct}")
+    return base64.b64decode(b64.encode())
+
+
+def ndarray_to_json(arr: n.ndarray):
+    """
+    Encode a numpy array a JSON-serializeable object.
+    """
+    return {
+        "$ndarray": {
+            "base64": base64.b64encode(arr.data).decode(),
+            "dtype": str(arr.dtype),
+            "shape": arr.shape,
+        }
+    }
+
+
+def ndarray_from_json(dct: Mapping[str, Any]):
+    """
+    Decode a serialized numpy array.
+    """
+    data = base64.b64decode(dct["$ndarray"]["base64"])
+    return n.frombuffer(data, dct["$ndarray"]["dtype"]).reshape(dct["$ndarray"]["shape"])

cryosparc/models/auth.py

@@ -0,0 +1,8 @@
+# THIS FILE IS AUTO-GENERATED, DO NOT EDIT DIRECTLY
+# SEE dev/api_generate_models.py
+from pydantic import BaseModel
+
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str

cryosparc/registry.py

@@ -0,0 +1,129 @@
+"""
+Model registration functions used by API client to determine how to interpret
+JSON responses. Used for either cryosparc-tools or cryosparc models.
+"""
+
+import re
+from collections.abc import Iterable
+from enum import Enum
+from inspect import isclass
+from types import ModuleType
+from typing import Dict, Optional, Type
+
+from pydantic import BaseModel
+
+from .stream import Streamable
+
+FINALIZED: bool = False
+REGISTERED_TYPED_DICTS: Dict[str, Type[dict]] = {}
+REGISTERED_ENUMS: Dict[str, Type[Enum]] = {}
+REGISTERED_MODEL_CLASSES: Dict[str, Type[BaseModel]] = {}
+REGISTERED_STREAM_CLASSES: Dict[str, Type[Streamable]] = {}
+
+
+def finalize():
+    """
+    Prevent registering additional types. Cannot be called twice.
+    """
+    global FINALIZED
+    check_finalized(False)
+    FINALIZED = True
+
+
+def check_finalized(finalized: bool = True):
+    """
+    Ensure the register has or hasn't been finalized. This is used in
+    special contexts such as cryosparcm icli or Jupyter Notebooks where
+    cryosparc-tools may be used alongside an API client.
+    """
+    assert FINALIZED is finalized, (
+        f"Cannot proceed because registry is {'finalized' if FINALIZED else 'not finalized'}. "
+        "This likely means that you're using both cryosparc-tools AND the "
+        "CryoSPARC API client from client/api_client.py. Please use either "
+        "`CryoSPARC` from tools or `APIClient` from cryosparc, but not both."
+    )
+
+
+def register_model(name, model_class: Type[BaseModel]):
+    check_finalized(False)
+    REGISTERED_MODEL_CLASSES[name] = model_class
+
+
+def register_typed_dict(name, typed_dict_class: Type[dict]):
+    check_finalized(False)
+    REGISTERED_TYPED_DICTS[name] = typed_dict_class
+
+
+def register_enum(name, enum_class: Type[Enum]):
+    check_finalized(False)
+    REGISTERED_ENUMS[name] = enum_class
+
+
+def register_model_module(mod: ModuleType):
+    for key, val in mod.__dict__.items():
+        if not re.match(r"^[A-Z]", key) or not isclass(val):
+            continue
+        if issubclass(val, BaseModel):
+            register_model(key, val)
+        if issubclass(val, dict):
+            register_typed_dict(key, val)
+        if issubclass(val, Enum):
+            register_enum(key, val)
+
+
+def model_for_ref(schema_ref: str) -> Optional[Type]:
+    """
+    Given a string with format either `#/components/schemas/X` or
+    `#/components/schemas/X_Y_`, looks up key X in `REGISTERED_MODEL_CLASSES``,
+    and return either X or X[Y] depending on whether the string includes the
+    final Y component.
+
+    Returns None if ref is not found.
+    """
+    import warnings
+
+    components = schema_ref.split("/")
+    if len(components) != 4 or components[0] != "#" or components[1] != "components" or components[2] != "schemas":
+        warnings.warn(f"Warning: Invalid schema reference {schema_ref}", stacklevel=2)
+        return
+
+    schema_name = components[3]
+    if "_" in schema_name:  # type var
+        generic, var, *_ = schema_name.split("_")
+        if generic in REGISTERED_MODEL_CLASSES and var in REGISTERED_MODEL_CLASSES:
+            return REGISTERED_MODEL_CLASSES[generic][REGISTERED_MODEL_CLASSES[var]]  # type: ignore
+    elif schema_name in REGISTERED_MODEL_CLASSES:
+        return REGISTERED_MODEL_CLASSES[schema_name]
+    elif schema_name in REGISTERED_TYPED_DICTS:
+        return REGISTERED_TYPED_DICTS[schema_name]
+    elif schema_name in REGISTERED_ENUMS:
+        return REGISTERED_ENUMS[schema_name]
+
+    warnings.warn(f"Warning: Unknown schema reference model {schema_ref}", stacklevel=2)
+
+
+def is_streamable_mime_type(mime: str):
+    return mime in REGISTERED_STREAM_CLASSES
+
+
+def register_stream_class(stream_class: Type[Streamable]):
+    mime = stream_class.media_type
+    assert mime not in REGISTERED_STREAM_CLASSES, (
+        f"Cannot register {stream_class}; "
+        f"stream class with mime-type {mime} is already registered "
+        f"({REGISTERED_STREAM_CLASSES[mime]})"
+    )
+    REGISTERED_STREAM_CLASSES[mime] = stream_class
+
+
+def get_stream_class(mime: str):
+    return REGISTERED_STREAM_CLASSES.get(mime)  # fails if mime-type not defined
+
+
+def streamable_mime_types():
+    return set(REGISTERED_STREAM_CLASSES.keys())
+
+
+def first_streamable_mime(strs: Iterable[str]) -> Optional[str]:
+    mimes = streamable_mime_types() & set(strs)
+    return mimes.pop() if len(mimes) > 0 else None

cryosparc/stream.py

@@ -107,13 +107,11 @@ async def read(self, n: Optional[int] = None):
 
 
 class Streamable(ABC):
-    @classmethod
-    def mime_type(cls) -> str:
-        """
-        Return the binary mime type to use in HTTP requests when streaming this
-        data e.g., "application/x-cryosparc-dataset"
-        """
-        return f"application/x-cryosparc-{cls.__name__.lower()}"
+    media_type = "application/octet-stream"
+    """
+    May override in subclasses to derive correct stream type, e.g.,
+    "application/x-cryosparc-dataset"
+    """
 
     @classmethod
     def api_schema(cls):
@@ -123,7 +121,7 @@ def api_schema(cls):
         """
         return {
             "description": f"A binary stream representing a CryoSPARC {cls.__name__}",
-            "content": {cls.mime_type(): {"schema": {"title": cls.__name__, "type": "string", "format": "binary"}}},
+            "content": {cls.media_type: {"schema": {"title": cls.__name__, "type": "string", "format": "binary"}}},
         }
 
     @classmethod

cryosparc/stream_registry.py

@@ -0,0 +1,4 @@
+from .dataset import Dataset
+from .registry import register_stream_class
+
+register_stream_class(Dataset)

pyproject.toml

@@ -22,6 +22,8 @@ classifiers = [
 license = { file = "LICENSE" }
 dependencies = [
     "numpy >= 1.17, < 3.0",
+    "httpx ~= 0.25",
+    "pydantic ~= 2.8",
     "typing-extensions >= 4.0",
 ]