Provide codec config schemas (#1)

* Prototype codec config schemas

* Derive codec config schemas for static codecs

* Test PyCodecClass schema polyfill

* Fix dict item iteration

* Check for Python __doc__ is None

* Improve PyCodec schema generation

* Refactor the PyCodec schema polyfill implementation

* Draft the docs extraction from a json schema, enums not yet supported

* Inline subschemas for static codecs

* Add enum schema and schema-derived PyCodec signature support

* Implement PyCodec[Class] downcast support

* Improvements to codec docs for enum configs

* Fix config parsing for SZ3 and ZFP

* Pin schemars to the latest alpha

* Refactor the schema parameter extraction

Author: juntyr

.gitignore

@@ -1,2 +1,4 @@
+/miniconda3
+
 /target
 /Cargo.lock

Cargo.toml

@@ -54,6 +54,7 @@ postcard = { version = "1.0", default-features = false }
 pyo3 = { version = "0.21", default-features = false }
 pythonize = { version = "0.21", default-features = false }
 rand = { version = "0.8", default-features = false }
+schemars = { version = "=1.0.0-alpha.9", default-features = false }
 serde = { version = "1.0", default-features = false }
 serde-transcode = { version = "1.1", default-features = false }
 serde_json = { version = "1.0", default-features = false }

codecs/bit-round/Cargo.toml

@@ -17,6 +17,7 @@ keywords = ["bit-rounding", "numcodecs", "compression", "encoding"]
 [dependencies]
 ndarray = { workspace = true }
 numcodecs = { workspace = true }
+schemars = { workspace = true, features = ["derive", "preserve_order"] }
 serde = { workspace = true, features = ["std", "derive"] }
 thiserror = { workspace = true }
 

codecs/bit-round/src/lib.rs

@@ -19,14 +19,24 @@
 
 use ndarray::{Array, ArrayBase, Data, Dimension};
 use numcodecs::{
-    serialize_codec_config_with_id, AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView,
-    AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
+    AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
+    Codec, StaticCodec, StaticCodecConfig,
 };
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-#[derive(Clone, Serialize, Deserialize)]
-/// Codec providing floating-point [`bit_round`]ing.
+#[derive(Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
+/// Codec providing floating-point bit rounding.
+///
+/// Drops the specified number of bits from the floating point mantissa,
+/// leaving an array that is more amenable to compression. The number of
+/// bits to keep should be determined by information analysis of the data
+/// to be compressed.
+///
+/// The approach is based on the paper by Klöwer et al. 2021
+/// (<https://www.nature.com/articles/s43588-021-00156-2>).
 pub struct BitRoundCodec {
     /// The number of bits of the mantissa to keep.
     ///
@@ -67,17 +77,19 @@ impl Codec for BitRoundCodec {
 
         Ok(decoded.assign(&encoded)?)
     }
-
-    fn get_config<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
-        serialize_codec_config_with_id(self, self, serializer)
-    }
 }
 
 impl StaticCodec for BitRoundCodec {
     const CODEC_ID: &'static str = "bit-round";
 
-    fn from_config<'de, D: Deserializer<'de>>(config: D) -> Result<Self, D::Error> {
-        Self::deserialize(config)
+    type Config<'de> = Self;
+
+    fn from_config(config: Self::Config<'_>) -> Self {
+        config
+    }
+
+    fn get_config(&self) -> StaticCodecConfig<Self> {
+        StaticCodecConfig::from(self)
     }
 }
 
@@ -104,15 +116,8 @@ pub enum BitRoundCodecError {
     },
 }
 
-/// Floating-point bit rounding.
-///
-/// Drops the specified number of bits from the floating point mantissa,
-/// leaving an array that is more amenable to compression. The number of
-/// bits to keep should be determined by information analysis of the data
-/// to be compressed.
-///
-/// The approach is based on the paper by Klöwer et al. 2021
-/// (<https://www.nature.com/articles/s43588-021-00156-2>).
+/// Floating-point bit rounding, which drops the specified number of bits from
+/// the floating point mantissa.
 ///
 /// See <https://github.com/milankl/BitInformation.jl> for the the original
 /// implementation in Julia.

codecs/identity/Cargo.toml

@@ -16,6 +16,7 @@ keywords = ["identity", "numcodecs", "compression", "encoding"]
 
 [dependencies]
 numcodecs = { workspace = true }
+schemars = { workspace = true, features = ["derive", "preserve_order"] }
 serde = { workspace = true, features = ["std", "derive"] }
 thiserror = { workspace = true }
 

codecs/identity/src/lib.rs

@@ -18,13 +18,15 @@
 //! Identity codec implementation for the [`numcodecs`] API.
 
 use numcodecs::{
-    serialize_codec_config_with_id, AnyArray, AnyArrayAssignError, AnyArrayView, AnyArrayViewMut,
-    AnyCowArray, Codec, StaticCodec,
+    AnyArray, AnyArrayAssignError, AnyArrayView, AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
+    StaticCodecConfig,
 };
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
 /// Identity codec which applies the identity function, i.e. passes through the
 /// input unchanged during encoding and decoding.
 pub struct IdentityCodec {
@@ -49,17 +51,19 @@ impl Codec for IdentityCodec {
     ) -> Result<(), Self::Error> {
         Ok(decoded.assign(&encoded)?)
     }
-
-    fn get_config<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
-        serialize_codec_config_with_id(self, self, serializer)
-    }
 }
 
 impl StaticCodec for IdentityCodec {
     const CODEC_ID: &'static str = "identity";
 
-    fn from_config<'de, D: Deserializer<'de>>(config: D) -> Result<Self, D::Error> {
-        Self::deserialize(config)
+    type Config<'de> = Self;
+
+    fn from_config(config: Self::Config<'_>) -> Self {
+        config
+    }
+
+    fn get_config(&self) -> StaticCodecConfig<Self> {
+        StaticCodecConfig::from(self)
     }
 }
 

codecs/linear-quantize/Cargo.toml

@@ -18,6 +18,7 @@ keywords = ["linear", "quantization", "numcodecs", "compression", "encoding"]
 ndarray = { workspace = true, features = ["std"] }
 numcodecs = { workspace = true }
 postcard = { workspace = true }
+schemars = { workspace = true, features = ["derive", "preserve_order"] }
 serde = { workspace = true, features = ["std", "derive"] }
 serde_repr = { workspace = true }
 thiserror = { workspace = true }

codecs/linear-quantize/src/lib.rs

@@ -23,15 +23,17 @@ use std::{borrow::Cow, fmt};
 
 use ndarray::{Array, Array1, ArrayBase, ArrayD, ArrayViewMutD, Data, Dimension, ShapeError};
 use numcodecs::{
-    serialize_codec_config_with_id, AnyArray, AnyArrayDType, AnyArrayView, AnyArrayViewMut,
-    AnyCowArray, Codec, StaticCodec,
+    AnyArray, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
+    StaticCodecConfig,
 };
-use serde::{de::DeserializeOwned, Deserialize, Deserializer, Serialize, Serializer};
+use schemars::{JsonSchema, JsonSchema_repr};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use serde_repr::{Deserialize_repr, Serialize_repr};
 use thiserror::Error;
 use twofloat::TwoFloat;
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
 /// Lossy codec to reduce the precision of floating point data.
 ///
 /// The data is quantized to unsigned integers of the best-fitting type.
@@ -44,7 +46,8 @@ pub struct LinearQuantizeCodec {
 }
 
 /// Data types which the [`LinearQuantizeCodec`] can quantize
-#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, JsonSchema)]
+#[schemars(extend("enum" = ["f32", "float32", "f64", "float64"]))]
 #[allow(missing_docs)]
 pub enum LinearQuantizeDType {
     #[serde(rename = "f32", alias = "float32")]
@@ -67,7 +70,7 @@ impl fmt::Display for LinearQuantizeDType {
 /// The binary `#[repr(u8)]` value of each variant is equivalent to the binary
 /// logarithm of the number of bins, i.e. the binary precision or the number of
 /// bits used.
-#[derive(Copy, Clone, Serialize_repr, Deserialize_repr)]
+#[derive(Copy, Clone, Serialize_repr, Deserialize_repr, JsonSchema_repr)]
 #[repr(u8)]
 #[rustfmt::skip]
 #[allow(missing_docs)]
@@ -382,17 +385,19 @@ impl Codec for LinearQuantizeCodec {
 
         Ok(())
     }
-
-    fn get_config<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
-        serialize_codec_config_with_id(self, self, serializer)
-    }
 }
 
 impl StaticCodec for LinearQuantizeCodec {
     const CODEC_ID: &'static str = "linear-quantize";
 
-    fn from_config<'de, D: Deserializer<'de>>(config: D) -> Result<Self, D::Error> {
-        Self::deserialize(config)
+    type Config<'de> = Self;
+
+    fn from_config(config: Self::Config<'_>) -> Self {
+        config
+    }
+
+    fn get_config(&self) -> StaticCodecConfig<Self> {
+        StaticCodecConfig::from(self)
     }
 }
 
@@ -714,7 +719,7 @@ impl Unsigned for u64 {
     const ZERO: Self = 0;
 }
 
-#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Serialize, Deserialize)]
 #[serde(bound = "")]
 struct CompressionHeader<'a, T: Float> {
     #[serde(borrow)]

codecs/log/Cargo.toml

@@ -17,6 +17,7 @@ keywords = ["log", "numcodecs", "compression", "encoding"]
 [dependencies]
 ndarray = { workspace = true }
 numcodecs = { workspace = true }
+schemars = { workspace = true, features = ["derive", "preserve_order"] }
 serde = { workspace = true, features = ["std", "derive"] }
 thiserror = { workspace = true }
 

codecs/log/src/lib.rs

@@ -19,13 +19,15 @@
 
 use ndarray::{Array, ArrayBase, ArrayView, ArrayViewMut, Data, Dimension};
 use numcodecs::{
-    serialize_codec_config_with_id, AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView,
-    AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
+    AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
+    Codec, StaticCodec, StaticCodecConfig,
 };
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
 /// Log codec which calculates `c = log(1+x)` on encoding and `d = exp(c)-1` on
 /// decoding.
 ///
@@ -76,17 +78,19 @@ impl Codec for LogCodec {
             (encoded, _decoded) => Err(LogCodecError::UnsupportedDtype(encoded.dtype())),
         }
     }
-
-    fn get_config<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
-        serialize_codec_config_with_id(self, self, serializer)
-    }
 }
 
 impl StaticCodec for LogCodec {
     const CODEC_ID: &'static str = "log";
 
-    fn from_config<'de, D: Deserializer<'de>>(config: D) -> Result<Self, D::Error> {
-        Self::deserialize(config)
+    type Config<'de> = Self;
+
+    fn from_config(config: Self::Config<'_>) -> Self {
+        config
+    }
+
+    fn get_config(&self) -> StaticCodecConfig<Self> {
+        StaticCodecConfig::from(self)
     }
 }