1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
use pyo3::{
    ffi::PyTypeObject,
    intern,
    prelude::*,
    sync::GILOnceCell,
    types::{DerefToPyAny, IntoPyDict, PyDict, PyType},
    PyTypeInfo,
};

#[allow(unused_imports)] // FIXME: use expect, only used in docs
use crate::PyCodecClassMethods;
use crate::{sealed::Sealed, PyCodecClass};

/// Represents a [`numcodecs.abc.Codec`] *instance* object.
///
/// The [`Bound<Codec>`] type implements the [`PyCodecMethods`] API.
///
/// [`numcodecs.abc.Codec`]: https://numcodecs.readthedocs.io/en/stable/abc.html#module-numcodecs.abc
#[repr(transparent)]
pub struct PyCodec {
    _codec: PyAny,
}

/// Methods implemented for [`PyCodec`]s.
pub trait PyCodecMethods<'py>: Sealed {
    /// Encodes the data in the buffer `buf` and returns the result.
    ///
    /// The input and output buffers be any objects supporting the
    /// [new-style buffer protocol].
    ///
    /// # Errors
    ///
    /// Errors if encoding the buffer fails.
    ///
    /// [new-style buffer protocol]: https://docs.python.org/3/c-api/buffer.html
    fn encode(&self, buf: Borrowed<'_, 'py, PyAny>) -> Result<Bound<'py, PyAny>, PyErr>;

    /// Decodes the data in the buffer `buf` and returns the result.
    ///
    /// The input and output buffers be any objects supporting the
    /// [new-style buffer protocol].
    ///
    /// If the optional output buffer `out` is provided, the decoded data is
    /// written into `out` and the `out` buffer is returned. Note that this
    /// buffer must be exactly the right size to store the decoded data.
    ///
    /// If the optional output buffer `out` is *not* provided, a new output
    /// buffer is allocated.
    ///
    /// # Errors
    ///
    /// Errors if decoding the buffer fails.
    ///
    /// [new-style buffer protocol]: https://docs.python.org/3/c-api/buffer.html
    fn decode(
        &self,
        buf: Borrowed<'_, 'py, PyAny>,
        out: Option<Borrowed<'_, 'py, PyAny>>,
    ) -> Result<Bound<'py, PyAny>, PyErr>;

    /// Returns a dictionary holding configuration parameters for this codec.
    ///
    /// The dict *must* include an `id` field with the
    /// [`PyCodecClassMethods::codec_id`].
    ///
    /// The dict *must* be compatible with JSON encoding.
    ///
    /// # Errors
    ///
    /// Errors if getting the codec configuration fails.
    fn get_config(&self) -> Result<Bound<'py, PyDict>, PyErr>;

    /// Returns the [`PyCodecClass`] of this codec.
    fn class(&self) -> Bound<'py, PyCodecClass>;
}

impl<'py> PyCodecMethods<'py> for Bound<'py, PyCodec> {
    fn encode(&self, buf: Borrowed<'_, 'py, PyAny>) -> Result<Bound<'py, PyAny>, PyErr> {
        let py = self.py();

        self.as_any().call_method1(intern!(py, "encode"), (buf,))
    }

    fn decode(
        &self,
        buf: Borrowed<'_, 'py, PyAny>,
        out: Option<Borrowed<'_, 'py, PyAny>>,
    ) -> Result<Bound<'py, PyAny>, PyErr> {
        let py = self.as_any().py();

        self.as_any().call_method(
            intern!(py, "decode"),
            (buf,),
            Some(&[(intern!(py, "out"), out)].into_py_dict_bound(py)),
        )
    }

    fn get_config(&self) -> Result<Bound<'py, PyDict>, PyErr> {
        let py = self.as_any().py();

        self.as_any()
            .call_method0(intern!(py, "get_config"))?
            .extract()
    }

    #[allow(clippy::expect_used)]
    fn class(&self) -> Bound<'py, PyCodecClass> {
        // extracting a codec guarantees that its class is a codec class
        self.as_any()
            .get_type()
            .extract()
            .expect("Codec's class must be a CodecClass")
    }
}

impl<'py> Sealed for Bound<'py, PyCodec> {}

#[doc(hidden)]
impl DerefToPyAny for PyCodec {}

#[doc(hidden)]
#[allow(unsafe_code)]
unsafe impl PyNativeType for PyCodec {
    type AsRefSource = Self;
}

#[doc(hidden)]
#[allow(unsafe_code)]
unsafe impl PyTypeInfo for PyCodec {
    const MODULE: Option<&'static str> = Some("numcodecs.abc");
    const NAME: &'static str = "Codec";

    #[inline]
    fn type_object_raw(py: Python) -> *mut PyTypeObject {
        static CODEC_TYPE: GILOnceCell<Py<PyType>> = GILOnceCell::new();

        let ty = CODEC_TYPE.get_or_try_init(py, || {
            py.import_bound(intern!(py, "numcodecs.abc"))?
                .getattr(intern!(py, "Codec"))?
                .extract()
        });
        #[allow(clippy::expect_used)]
        let ty = ty.expect("failed to access the `numpy.abc.Codec` type object");

        ty.bind(py).as_type_ptr()
    }
}