quantity/
python.rs

1use super::{Angle, Quantity, SIUnit};
2use crate::fmt::PrintUnit;
3#[cfg(feature = "ndarray")]
4use ndarray::{Array, Dimension};
5#[cfg(feature = "ndarray")]
6use numpy::{IntoPyArray, PyReadonlyArray};
7use pyo3::{exceptions::PyValueError, prelude::*};
8use std::{marker::PhantomData, sync::LazyLock};
9use typenum::Integer;
10
11static SIOBJECT: LazyLock<PyObject> = LazyLock::new(|| {
12    Python::with_gil(|py| {
13        PyModule::import(py, "si_units")
14            .unwrap()
15            .getattr("SIObject")
16            .unwrap()
17            .unbind()
18    })
19});
20
21impl<
22        'py,
23        T: Integer,
24        L: Integer,
25        M: Integer,
26        I: Integer,
27        THETA: Integer,
28        N: Integer,
29        J: Integer,
30    > IntoPyObject<'py> for Quantity<f64, SIUnit<T, L, M, I, THETA, N, J>>
31{
32    type Target = PyAny;
33    type Output = Bound<'py, PyAny>;
34    type Error = PyErr;
35
36    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
37        let unit = [L::I8, M::I8, T::I8, I::I8, N::I8, THETA::I8, J::I8];
38        SIOBJECT.bind(py).call1((self.0, unit))
39    }
40}
41
42#[cfg(feature = "ndarray")]
43impl<
44        'py,
45        T: Integer,
46        L: Integer,
47        M: Integer,
48        I: Integer,
49        THETA: Integer,
50        N: Integer,
51        J: Integer,
52        D: Dimension,
53    > IntoPyObject<'py> for Quantity<Array<f64, D>, SIUnit<T, L, M, I, THETA, N, J>>
54{
55    type Target = PyAny;
56    type Output = Bound<'py, PyAny>;
57    type Error = PyErr;
58
59    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
60        let unit = [L::I8, M::I8, T::I8, I::I8, N::I8, THETA::I8, J::I8];
61        let value = self.0.into_pyarray(py).into_any();
62        SIOBJECT.bind(py).call1((value, unit))
63    }
64}
65
66impl<
67        'py,
68        T: Integer,
69        L: Integer,
70        M: Integer,
71        I: Integer,
72        THETA: Integer,
73        N: Integer,
74        J: Integer,
75    > FromPyObject<'py> for Quantity<f64, SIUnit<T, L, M, I, THETA, N, J>>
76where
77    Self: PrintUnit,
78{
79    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
80        let Ok((value, unit_from)) = ob
81            .call_method0("__getnewargs__")
82            .and_then(|raw| raw.extract::<(f64, [i8; 7])>())
83        else {
84            return Err(PyErr::new::<PyValueError, _>(format!(
85                "Missing units! Expected {}, got {}.",
86                Self::UNIT,
87                ob.call_method0("__repr__")?
88            )));
89        };
90        let unit_into = [L::I8, M::I8, T::I8, I::I8, N::I8, THETA::I8, J::I8];
91        if unit_into == unit_from {
92            Ok(Quantity(value, PhantomData))
93        } else {
94            Err(PyErr::new::<PyValueError, _>(format!(
95                "Wrong units! Expected {}, got {}.",
96                Self::UNIT,
97                ob.call_method0("__repr__")?
98            )))
99        }
100    }
101}
102
103#[cfg(feature = "ndarray")]
104impl<
105        'py,
106        T: Integer,
107        L: Integer,
108        M: Integer,
109        I: Integer,
110        THETA: Integer,
111        N: Integer,
112        J: Integer,
113        D: Dimension,
114    > FromPyObject<'py> for Quantity<Array<f64, D>, SIUnit<T, L, M, I, THETA, N, J>>
115where
116    Self: PrintUnit,
117{
118    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
119        let Ok((value, unit_from)) = ob
120            .call_method0("__getnewargs__")
121            .and_then(|raw| raw.extract::<(PyReadonlyArray<f64, D>, [i8; 7])>())
122        else {
123            return Err(PyErr::new::<PyValueError, _>(format!(
124                "Missing units! Expected {}, got {}.",
125                Self::UNIT,
126                ob.call_method0("__repr__")?
127            )));
128        };
129        let value = value.as_array().to_owned();
130        let unit_into = [L::I8, M::I8, T::I8, I::I8, N::I8, THETA::I8, J::I8];
131        if unit_into == unit_from {
132            Ok(Quantity(value, PhantomData))
133        } else {
134            Err(PyErr::new::<PyValueError, _>(format!(
135                "Wrong units! Expected {}, got {}.",
136                Self::UNIT,
137                ob.call_method0("__repr__")?
138            )))
139        }
140    }
141}
142
143static ANGLE: LazyLock<PyObject> = LazyLock::new(|| {
144    Python::with_gil(|py| {
145        PyModule::import(py, "si_units")
146            .unwrap()
147            .getattr("Angle")
148            .unwrap()
149            .unbind()
150    })
151});
152
153impl<'py> IntoPyObject<'py> for Angle {
154    type Target = PyAny;
155    type Output = Bound<'py, PyAny>;
156    type Error = PyErr;
157    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
158        ANGLE.bind(py).call1((self.0,))
159    }
160}
161
162impl<'py> FromPyObject<'py> for Angle {
163    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
164        let Ok(value) = ob
165            .call_method0("__getnewargs__")
166            .and_then(|raw| raw.extract::<f64>())
167        else {
168            return Err(PyErr::new::<PyValueError, _>(format!(
169                "Missing units! Expected angle, got {}.",
170                ob.call_method0("__repr__")?
171            )));
172        };
173        Ok(Quantity(value, PhantomData))
174    }
175}