1use numcodecs_python::{PyCodec, PyCodecAdapter, PyCodecClass};
2use numcodecs_wasm_host_reproducible::{ReproducibleWasmCodec, ReproducibleWasmCodecType};
3use pyo3::{exceptions::PyTypeError, prelude::*};
4
5use ::memfd as _;
7
8mod engine;
9
10use engine::{Engine, default_engine};
11
12#[pymodule]
13#[pyo3(name = "_wasm")]
14fn wasm<'py>(py: Python<'py>, module: &Bound<'py, PyModule>) -> Result<(), PyErr> {
15 let logger = pyo3_log::Logger::new(py, pyo3_log::Caching::Nothing)?;
16 logger
17 .install()
18 .map_err(|err| pyo3_error::PyErrChain::new(py, err))?;
19
20 module.add_function(wrap_pyfunction!(create_codec_class, module)?)?;
21 module.add_function(wrap_pyfunction!(read_codec_instruction_counter, module)?)?;
22
23 Ok(())
24}
25
26#[pyfunction]
27#[pyo3(name = "_create_codec_class")]
28fn create_codec_class<'py>(
29 py: Python<'py>,
30 module: &Bound<'py, PyModule>,
31 wasm: Vec<u8>,
32) -> Result<Bound<'py, PyCodecClass>, PyErr> {
33 let engine = default_engine(py)?;
34
35 let codec_ty = ReproducibleWasmCodecType::new(engine, wasm)
36 .map_err(|err| pyo3_error::PyErrChain::new(py, err))?;
37
38 let codec_class = numcodecs_python::export_codec_class(py, codec_ty, module.as_borrowed())?;
39
40 Ok(codec_class)
41}
42
43#[pyfunction]
44#[pyo3(name = "_read_codec_instruction_counter")]
45fn read_codec_instruction_counter<'py>(
46 py: Python<'py>,
47 codec: &Bound<'py, PyCodec>,
48) -> Result<u64, PyErr> {
49 let Some(instruction_counter) =
50 PyCodecAdapter::with_downcast(py, codec, |codec: &ReproducibleWasmCodec<Engine>| {
51 codec.instruction_counter()
52 })
53 .transpose()
54 .map_err(|err| pyo3_error::PyErrChain::new(py, err))?
55 else {
56 return Err(PyTypeError::new_err(
57 "`codec` is not a wasm codec, only wasm codecs have instruction counts",
58 ));
59 };
60
61 Ok(instruction_counter.0)
62}