numpy/
strings.rs

1//! Types to support arrays of [ASCII][ascii] and [UCS4][ucs4] strings
2//!
3//! [ascii]: https://numpy.org/doc/stable/reference/c-api/dtype.html#c.NPY_STRING
4//! [ucs4]: https://numpy.org/doc/stable/reference/c-api/dtype.html#c.NPY_UNICODE
5
6use std::collections::hash_map::Entry;
7use std::ffi::c_char;
8use std::fmt;
9use std::mem::size_of;
10use std::str;
11use std::sync::Mutex;
12
13use pyo3::sync::MutexExt;
14use pyo3::{
15    Bound, Py, Python,
16    ffi::{Py_UCS1, Py_UCS4},
17};
18use rustc_hash::FxHashMap;
19
20use crate::dtype::{Element, PyArrayDescr, PyArrayDescrMethods, clone_methods_impl};
21use crate::npyffi::{_PyDataType_GET_ITEM_DATA, NPY_TYPES, PyDataType_SET_ELSIZE};
22
23/// A newtype wrapper around [`[u8; N]`][Py_UCS1] to handle [`byte` scalars][numpy-bytes] while satisfying coherence.
24///
25/// Note that when creating arrays of ASCII strings without an explicit `dtype`,
26/// NumPy will automatically determine the smallest possible array length at runtime.
27///
28/// For example,
29///
30/// ```python
31/// array = numpy.array([b"foo", b"bar", b"foobar"])
32/// ```
33///
34/// yields `S6` for `array.dtype`.
35///
36/// On the Rust side however, the length `N` of `PyFixedString<N>` must always be given
37/// explicitly and as a compile-time constant. For this work reliably, the Python code
38/// should set the `dtype` explicitly, e.g.
39///
40/// ```python
41/// numpy.array([b"foo", b"bar", b"foobar"], dtype='S12')
42/// ```
43///
44/// always matching `PyArray1<PyFixedString<12>>`.
45///
46/// # Example
47///
48/// ```rust
49/// # use pyo3::Python;
50/// use numpy::{PyArray1, PyUntypedArrayMethods, PyFixedString};
51///
52/// # Python::attach(|py| {
53/// let array = PyArray1::<PyFixedString<3>>::from_vec(py, vec![[b'f', b'o', b'o'].into()]);
54///
55/// assert!(array.dtype().to_string().contains("S3"));
56/// # });
57/// ```
58///
59/// [numpy-bytes]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_
60#[repr(transparent)]
61#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
62pub struct PyFixedString<const N: usize>(pub [Py_UCS1; N]);
63
64impl<const N: usize> fmt::Display for PyFixedString<N> {
65    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
66        fmt.write_str(str::from_utf8(&self.0).unwrap().trim_end_matches('\0'))
67    }
68}
69
70impl<const N: usize> From<[Py_UCS1; N]> for PyFixedString<N> {
71    fn from(val: [Py_UCS1; N]) -> Self {
72        Self(val)
73    }
74}
75
76unsafe impl<const N: usize> Element for PyFixedString<N> {
77    const IS_COPY: bool = true;
78
79    fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr> {
80        static DTYPES: TypeDescriptors = TypeDescriptors::new();
81
82        unsafe { DTYPES.from_size(py, NPY_TYPES::NPY_STRING, b'|' as _, size_of::<Self>()) }
83    }
84
85    clone_methods_impl!(Self);
86}
87
88/// A newtype wrapper around [`[PyUCS4; N]`][Py_UCS4] to handle [`str_` scalars][numpy-str] while satisfying coherence.
89///
90/// Note that when creating arrays of Unicode strings without an explicit `dtype`,
91/// NumPy will automatically determine the smallest possible array length at runtime.
92///
93/// For example,
94///
95/// ```python
96/// numpy.array(["foo🐍", "bar🦀", "foobar"])
97/// ```
98///
99/// yields `U6` for `array.dtype`.
100///
101/// On the Rust side however, the length `N` of `PyFixedUnicode<N>` must always be given
102/// explicitly and as a compile-time constant. For this work reliably, the Python code
103/// should set the `dtype` explicitly, e.g.
104///
105/// ```python
106/// numpy.array(["foo🐍", "bar🦀", "foobar"], dtype='U12')
107/// ```
108///
109/// always matching `PyArray1<PyFixedUnicode<12>>`.
110///
111/// # Example
112///
113/// ```rust
114/// # use pyo3::Python;
115/// use numpy::{PyArray1, PyUntypedArrayMethods, PyFixedUnicode};
116///
117/// # Python::attach(|py| {
118/// let array = PyArray1::<PyFixedUnicode<3>>::from_vec(py, vec![[b'b' as _, b'a' as _, b'r' as _].into()]);
119///
120/// assert!(array.dtype().to_string().contains("U3"));
121/// # });
122/// ```
123///
124/// [numpy-str]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.str_
125#[repr(transparent)]
126#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
127pub struct PyFixedUnicode<const N: usize>(pub [Py_UCS4; N]);
128
129impl<const N: usize> fmt::Display for PyFixedUnicode<N> {
130    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
131        for character in self.0 {
132            if character == 0 {
133                break;
134            }
135
136            write!(fmt, "{}", char::from_u32(character).unwrap())?;
137        }
138
139        Ok(())
140    }
141}
142
143impl<const N: usize> From<[Py_UCS4; N]> for PyFixedUnicode<N> {
144    fn from(val: [Py_UCS4; N]) -> Self {
145        Self(val)
146    }
147}
148
149unsafe impl<const N: usize> Element for PyFixedUnicode<N> {
150    const IS_COPY: bool = true;
151
152    fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr> {
153        static DTYPES: TypeDescriptors = TypeDescriptors::new();
154
155        unsafe { DTYPES.from_size(py, NPY_TYPES::NPY_UNICODE, b'=' as _, size_of::<Self>()) }
156    }
157
158    clone_methods_impl!(Self);
159}
160
161struct TypeDescriptors {
162    dtypes: Mutex<Option<FxHashMap<usize, Py<PyArrayDescr>>>>,
163}
164
165impl TypeDescriptors {
166    const fn new() -> Self {
167        Self {
168            dtypes: Mutex::new(None),
169        }
170    }
171
172    /// `npy_type` must be either `NPY_STRING` or `NPY_UNICODE` with matching `byteorder` and `size`
173    #[allow(clippy::wrong_self_convention)]
174    unsafe fn from_size<'py>(
175        &self,
176        py: Python<'py>,
177        npy_type: NPY_TYPES,
178        byteorder: c_char,
179        size: usize,
180    ) -> Bound<'py, PyArrayDescr> {
181        let mut dtypes = self
182            .dtypes
183            .lock_py_attached(py)
184            .expect("dtype cache poisoned");
185
186        let dtype = match dtypes.get_or_insert_with(Default::default).entry(size) {
187            Entry::Occupied(entry) => entry.into_mut(),
188            Entry::Vacant(entry) => {
189                let dtype = PyArrayDescr::new_from_npy_type(py, npy_type);
190
191                let descr = dtype.as_dtype_ptr();
192                unsafe {
193                    PyDataType_SET_ELSIZE(py, descr, size.try_into().unwrap());
194                    (*_PyDataType_GET_ITEM_DATA(descr)).byteorder = byteorder;
195                }
196
197                entry.insert(dtype.into())
198            }
199        };
200
201        dtype.bind(py).to_owned()
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    #[allow(clippy::byte_char_slices)]
211    fn format_fixed_string() {
212        assert_eq!(
213            PyFixedString([b'f', b'o', b'o', 0, 0, 0]).to_string(),
214            "foo"
215        );
216        assert_eq!(PyFixedString(*b"foobar").to_string(), "foobar");
217    }
218
219    #[test]
220    fn format_fixed_unicode() {
221        assert_eq!(
222            PyFixedUnicode([b'f' as _, b'o' as _, b'o' as _, 0, 0, 0]).to_string(),
223            "foo"
224        );
225        assert_eq!(
226            PyFixedUnicode([0x1F980, 0x1F40D, 0, 0, 0, 0]).to_string(),
227            "🦀🐍"
228        );
229        assert_eq!(
230            PyFixedUnicode([
231                b'f' as _, b'o' as _, b'o' as _, b'b' as _, b'a' as _, b'r' as _
232            ])
233            .to_string(),
234            "foobar"
235        );
236    }
237}
numpy/strings.rs

numpy/
strings.rs
