numcodecs_python/
registry.rs

1use numcodecs::{DynCodec, ErasedDynCodec};
2use numcodecs_registry::Registry;
3use pyo3::{prelude::*, sync::PyOnceLock, types::PyDict};
4use pythonize::Pythonizer;
5use serde::Deserializer;
6use serde_transcode::transcode;
7
8#[expect(unused_imports)] // FIXME: use expect, only used in docs
9use crate::PyCodecClassMethods;
10use crate::{PyCodec, PyCodecAdapter, PyCodecClass};
11
12/// Dynamic registry of codec classes.
13pub struct PyCodecRegistry {
14    _private: (),
15}
16
17impl PyCodecRegistry {
18    /// Instantiate a codec from a configuration dictionary.
19    ///
20    /// The config *must* include the `id` field with the
21    /// [`PyCodecClassMethods::codec_id`].
22    ///
23    /// # Errors
24    ///
25    /// Errors if no codec with a matching `id` has been registered, or if
26    /// constructing the codec fails.
27    pub fn get_codec<'py>(config: Borrowed<'_, 'py, PyDict>) -> Result<Bound<'py, PyCodec>, PyErr> {
28        static GET_CODEC: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
29
30        let py = config.py();
31
32        let get_codec = GET_CODEC.import(py, "numcodecs.registry", "get_codec")?;
33
34        get_codec.call1((config,))?.extract()
35    }
36
37    /// Register a codec class.
38    ///
39    /// If the `codec_id` is provided, it is used instead of
40    /// [`PyCodecClassMethods::codec_id`].
41    ///
42    /// This function maintains a mapping from codec identifiers to codec
43    /// classes. When a codec class is registered, it will replace any class
44    /// previously registered under the same codec identifier, if present.
45    ///
46    /// # Errors
47    ///
48    /// Errors if registering the codec class fails.
49    pub fn register_codec(
50        class: Borrowed<PyCodecClass>,
51        codec_id: Option<&str>,
52    ) -> Result<(), PyErr> {
53        static REGISTER_CODEC: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
54
55        let py = class.py();
56
57        let register_codec = REGISTER_CODEC.import(py, "numcodecs.registry", "register_codec")?;
58
59        register_codec.call1((class, codec_id))?;
60
61        Ok(())
62    }
63}
64
65/// Handle for [`PyCodecRegistry`].
66pub struct PyCodecRegistryHandle;
67
68impl Registry for PyCodecRegistryHandle {
69    type Error = PyErr;
70
71    fn get_codec<'de, D: Deserializer<'de>>(
72        &self,
73        config: D,
74    ) -> Result<ErasedDynCodec, Self::Error> {
75        Python::attach(|py| {
76            let config = transcode(config, Pythonizer::new(py))?;
77            let config: Bound<PyDict> = config.extract()?;
78
79            let codec = PyCodecRegistry::get_codec(config.as_borrowed())?;
80            let codec = PyCodecAdapter::from_codec(codec)?;
81
82            Ok(ErasedDynCodec::new(codec))
83        })
84    }
85
86    fn get_codec_typed<'de, T: DynCodec, D: Deserializer<'de>>(
87        &self,
88        config: D,
89    ) -> Result<Option<T>, Self::Error> {
90        Python::attach(|py| {
91            let config = transcode(config, Pythonizer::new(py))?;
92            let config: Bound<PyDict> = config.extract()?;
93
94            let codec = PyCodecRegistry::get_codec(config.as_borrowed())?;
95            // clone is necessary since we cannot move out of a PyCodec
96            let codec = PyCodecAdapter::with_downcast(py, &codec, |codec: &T| codec.clone());
97
98            Ok(codec)
99        })
100    }
101}
numcodecs_python/registry.rs

numcodecs_python/
registry.rs
