Add datetime module for handling NumPy's datetime64 and timedelta64 types.

Author: adamreichold

CHANGELOG.md

@@ -4,6 +4,7 @@
   - Add dynamic borrow checking to safely construct references into the interior of NumPy arrays. ([#274](https://github.com/PyO3/rust-numpy/pull/274))
     - The deprecated iterator builders `NpySingleIterBuilder::{readonly,readwrite}` and `NpyMultiIterBuilder::add_{readonly,readwrite}` now take referencces to `PyReadonlyArray` and `PyReadwriteArray` instead of consuming them.
     - The destructive `PyArray::resize` method is now unsafe if used without an instance of `PyReadwriteArray`. ([#302](https://github.com/PyO3/rust-numpy/pull/302))
+  - Add support for `datetime64` and `timedelta64` element types via the `datetime` module. ([#308](https://github.com/PyO3/rust-numpy/pull/308))
   - Add support for IEEE 754-2008 16-bit floating point numbers via an optional dependency on the `half` crate. ([#314](https://github.com/PyO3/rust-numpy/pull/314))
   - The `inner`, `dot` and `einsum` functions can also return a scalar instead of a zero-dimensional array to match NumPy's types ([#285](https://github.com/PyO3/rust-numpy/pull/285))
   - The `PyArray::resize` function supports n-dimensional contiguous arrays. ([#312](https://github.com/PyO3/rust-numpy/pull/312))

examples/simple/src/lib.rs

@@ -1,6 +1,8 @@
-use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
+use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD, Zip};
 use numpy::{
-    Complex64, IntoPyArray, PyArray1, PyArrayDyn, PyReadonlyArrayDyn, PyReadwriteArrayDyn,
+    datetime::{units, Timedelta},
+    Complex64, IntoPyArray, PyArray1, PyArrayDyn, PyReadonlyArray1, PyReadonlyArrayDyn,
+    PyReadwriteArray1, PyReadwriteArrayDyn,
 };
 use pyo3::{
     pymodule,
@@ -70,5 +72,17 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
         x.readonly().as_array().sum()
     }
 
+    // example using timedelta64 array
+    #[pyfn(m)]
+    fn add_minutes_to_seconds(
+        mut x: PyReadwriteArray1<Timedelta<units::Seconds>>,
+        y: PyReadonlyArray1<Timedelta<units::Minutes>>,
+    ) {
+        #[allow(deprecated)]
+        Zip::from(x.as_array_mut())
+            .and(y.as_array())
+            .apply(|x, y| *x = (i64::from(*x) + 60 * i64::from(*y)).into());
+    }
+
     Ok(())
 }

examples/simple/tests/test_ext.py

@@ -1,5 +1,5 @@
 import numpy as np
-from rust_ext import axpy, conj, mult, extract
+from rust_ext import axpy, conj, mult, extract, add_minutes_to_seconds
 
 
 def test_axpy():
@@ -24,3 +24,12 @@ def test_extract():
     x = np.arange(5.0)
     d = {"x": x}
     np.testing.assert_almost_equal(extract(d), 10.0)
+
+
+def test_add_minutes_to_seconds():
+    x = np.array([10, 20, 30], dtype="timedelta64[s]")
+    y = np.array([1, 2, 3], dtype="timedelta64[m]")
+
+    add_minutes_to_seconds(x, y)
+
+    assert np.all(x == np.array([70, 140, 210], dtype="timedelta64[s]"))

src/datetime.rs

@@ -0,0 +1,309 @@
+//! Support datetimes and timedeltas
+//!
+//! [The corresponding section][datetime] of the NumPy documentation contains more information.
+//!
+//! # Example
+//!
+//! ```
+//! use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1};
+//! use pyo3::Python;
+//! # use pyo3::types::PyDict;
+//!
+//! Python::with_gil(|py| {
+//! #    let locals = py
+//! #        .eval("{ 'np': __import__('numpy') }", None, None)
+//! #        .unwrap()
+//! #        .downcast::<PyDict>()
+//! #        .unwrap();
+//! #
+//!     let array = py
+//!         .eval(
+//!             "np.array([np.datetime64('2017-04-21')])",
+//!             None,
+//!             Some(locals),
+//!         )
+//!         .unwrap()
+//!         .downcast::<PyArray1<Datetime<units::Days>>>()
+//!         .unwrap();
+//!
+//!     assert_eq!(
+//!         array.get_owned(0).unwrap(),
+//!         Datetime::<units::Days>::from(17_277)
+//!     );
+//!
+//!     let array = py
+//!         .eval(
+//!             "np.array([np.datetime64('2022-03-29')]) - np.array([np.datetime64('2017-04-21')])",
+//!             None,
+//!             Some(locals),
+//!         )
+//!         .unwrap()
+//!         .downcast::<PyArray1<Timedelta<units::Days>>>()
+//!         .unwrap();
+//!
+//!     assert_eq!(
+//!         array.get_owned(0).unwrap(),
+//!         Timedelta::<units::Days>::from(1_803)
+//!     );
+//! });
+//! ```
+//!
+//! [datetime]: https://numpy.org/doc/stable/reference/arrays.datetime.html
+
+use std::fmt;
+use std::hash::Hash;
+use std::marker::PhantomData;
+
+use pyo3::Python;
+
+use crate::dtype::{Element, PyArrayDescr};
+use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};
+
+/// Represents the [datetime units][datetime-units] supported by NumPy
+///
+/// [datetime-units]: https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
+pub trait Unit: Send + Sync + Clone + Copy + PartialEq + Eq + Hash + PartialOrd + Ord {
+    /// The matching NumPy [datetime unit code][NPY_DATETIMEUNIT]
+    ///
+    /// [NPY_DATETIMEUNIT]: https://github.com/numpy/numpy/blob/4c60b3263ac50e5e72f6a909e156314fc3c9cba0/numpy/core/include/numpy/ndarraytypes.h#L276
+    const UNIT: NPY_DATETIMEUNIT;
+
+    /// The abbrevation used for debug formatting
+    const ABBREV: &'static str;
+}
+
+macro_rules! define_units {
+    ($($(#[$meta:meta])* $struct:ident => $unit:ident $abbrev:literal,)+) => {
+        $(
+
+        $(#[$meta])*
+        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
+        pub struct $struct;
+
+        impl Unit for $struct {
+            const UNIT: NPY_DATETIMEUNIT = NPY_DATETIMEUNIT::$unit;
+
+            const ABBREV: &'static str = $abbrev;
+        }
+
+        )+
+    };
+}
+
+/// Predefined implementors of the [`Unit`] trait
+pub mod units {
+    use super::*;
+
+    define_units!(
+        #[doc = "Years, i.e. 12 months"]
+        Years => NPY_FR_Y "a",
+        #[doc = "Months, i.e. 30 days"]
+        Months => NPY_FR_M "mo",
+        #[doc = "Weeks, i.e. 7 days"]
+        Weeks => NPY_FR_W "w",
+        #[doc = "Days, i.e. 24 hours"]
+        Days => NPY_FR_D "d",
+        #[doc = "Hours, i.e. 60 minutes"]
+        Hours => NPY_FR_h "h",
+        #[doc = "Minutes, i.e. 60 seconds"]
+        Minutes => NPY_FR_m "min",
+        #[doc = "Seconds"]
+        Seconds => NPY_FR_s "s",
+        #[doc = "Milliseconds, i.e. 10^-3 seconds"]
+        Milliseconds => NPY_FR_ms "ms",
+        #[doc = "Microseconds, i.e. 10^-6 seconds"]
+        Microseconds => NPY_FR_us "µs",
+        #[doc = "Nanoseconds, i.e. 10^-9 seconds"]
+        Nanoseconds => NPY_FR_ns "ns",
+        #[doc = "Picoseconds, i.e. 10^-12 seconds"]
+        Picoseconds => NPY_FR_ps "ps",
+        #[doc = "Femtoseconds, i.e. 10^-15 seconds"]
+        Femtoseconds => NPY_FR_fs "fs",
+        #[doc = "Attoseconds, i.e. 10^-18 seconds"]
+        Attoseconds => NPY_FR_as "as",
+    );
+}
+
+/// Corresponds to the [`datetime64`][scalars-datetime64] scalar type
+///
+/// [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
+#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[repr(transparent)]
+pub struct Datetime<U: Unit>(i64, PhantomData<U>);
+
+impl<U: Unit> From<i64> for Datetime<U> {
+    fn from(val: i64) -> Self {
+        Self(val, PhantomData)
+    }
+}
+
+impl<U: Unit> From<Datetime<U>> for i64 {
+    fn from(val: Datetime<U>) -> Self {
+        val.0
+    }
+}
+
+unsafe impl<U: Unit> Element for Datetime<U> {
+    const IS_COPY: bool = true;
+
+    fn get_dtype(py: Python) -> &PyArrayDescr {
+        // FIXME(adamreichold): Memoize these via the Unit trait
+        let dtype = PyArrayDescr::new_from_npy_type(py, NPY_TYPES::NPY_DATETIME);
+
+        unsafe {
+            set_unit(dtype, U::UNIT);
+        }
+
+        dtype
+    }
+}
+
+impl<U: Unit> fmt::Debug for Datetime<U> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Datetime({} {})", self.0, U::ABBREV)
+    }
+}
+
+/// Corresponds to the [`timedelta64`][scalars-datetime64] scalar type
+///
+/// [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64
+#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[repr(transparent)]
+pub struct Timedelta<U: Unit>(i64, PhantomData<U>);
+
+impl<U: Unit> From<i64> for Timedelta<U> {
+    fn from(val: i64) -> Self {
+        Self(val, PhantomData)
+    }
+}
+
+impl<U: Unit> From<Timedelta<U>> for i64 {
+    fn from(val: Timedelta<U>) -> Self {
+        val.0
+    }
+}
+
+unsafe impl<U: Unit> Element for Timedelta<U> {
+    const IS_COPY: bool = true;
+
+    fn get_dtype(py: Python) -> &PyArrayDescr {
+        // FIXME(adamreichold): Memoize these via the Unit trait
+        let dtype = PyArrayDescr::new_from_npy_type(py, NPY_TYPES::NPY_TIMEDELTA);
+
+        unsafe {
+            set_unit(dtype, U::UNIT);
+        }
+
+        dtype
+    }
+}
+
+impl<U: Unit> fmt::Debug for Timedelta<U> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Timedelta({} {})", self.0, U::ABBREV)
+    }
+}
+
+unsafe fn set_unit(dtype: &PyArrayDescr, unit: NPY_DATETIMEUNIT) {
+    let metadata = &mut *((*dtype.as_dtype_ptr()).c_metadata as *mut PyArray_DatetimeDTypeMetaData);
+
+    metadata.meta.base = unit;
+    metadata.meta.num = 1;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use pyo3::{
+        py_run,
+        types::{PyDict, PyModule},
+    };
+
+    use crate::array::PyArray1;
+
+    #[test]
+    fn from_python_to_rust() {
+        Python::with_gil(|py| {
+            let locals = py
+                .eval("{ 'np': __import__('numpy') }", None, None)
+                .unwrap()
+                .downcast::<PyDict>()
+                .unwrap();
+
+            let array = py
+                .eval(
+                    "np.array([np.datetime64('1970-01-01')])",
+                    None,
+                    Some(locals),
+                )
+                .unwrap()
+                .downcast::<PyArray1<Datetime<units::Days>>>()
+                .unwrap();
+
+            let value: i64 = array.get_owned(0).unwrap().into();
+            assert_eq!(value, 0);
+        });
+    }
+
+    #[test]
+    fn from_rust_to_python() {
+        Python::with_gil(|py| {
+            let array = PyArray1::<Timedelta<units::Minutes>>::zeros(py, 1, false);
+
+            *array.readwrite().get_mut(0).unwrap() = Timedelta::<units::Minutes>::from(5);
+
+            let np = py
+                .eval("__import__('numpy')", None, None)
+                .unwrap()
+                .downcast::<PyModule>()
+                .unwrap();
+
+            py_run!(py, array np, "assert array.dtype == np.dtype('timedelta64[m]')");
+            py_run!(py, array np, "assert array[0] == np.timedelta64(5, 'm')");
+        });
+    }
+
+    #[test]
+    fn debug_formatting() {
+        assert_eq!(
+            format!("{:?}", Datetime::<units::Days>::from(28)),
+            "Datetime(28 d)"
+        );
+
+        assert_eq!(
+            format!("{:?}", Timedelta::<units::Milliseconds>::from(160)),
+            "Timedelta(160 ms)"
+        );
+    }
+
+    #[test]
+    fn unit_conversion() {
+        #[track_caller]
+        fn convert<S: Unit, D: Unit>(py: Python<'_>, expected_value: i64) {
+            let array = PyArray1::<Timedelta<S>>::from_slice(py, &[Timedelta::<S>::from(1)]);
+            let array = array.cast::<Timedelta<D>>(false).unwrap();
+
+            let value: i64 = array.get_owned(0).unwrap().into();
+            assert_eq!(value, expected_value);
+        }
+
+        Python::with_gil(|py| {
+            convert::<units::Years, units::Days>(py, (97 + 400 * 365) / 400);
+            convert::<units::Months, units::Days>(py, (97 + 400 * 365) / 400 / 12);
+
+            convert::<units::Weeks, units::Seconds>(py, 7 * 24 * 60 * 60);
+            convert::<units::Days, units::Seconds>(py, 24 * 60 * 60);
+            convert::<units::Hours, units::Seconds>(py, 60 * 60);
+            convert::<units::Minutes, units::Seconds>(py, 60);
+
+            convert::<units::Seconds, units::Milliseconds>(py, 1_000);
+            convert::<units::Seconds, units::Microseconds>(py, 1_000_000);
+            convert::<units::Seconds, units::Nanoseconds>(py, 1_000_000_000);
+            convert::<units::Seconds, units::Picoseconds>(py, 1_000_000_000_000);
+            convert::<units::Seconds, units::Femtoseconds>(py, 1_000_000_000_000_000);
+
+            convert::<units::Femtoseconds, units::Attoseconds>(py, 1_000);
+        });
+    }
+}

src/dtype.rs

@@ -127,13 +127,20 @@ impl PyArrayDescr {
         }
     }
 
-    pub(crate) fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
+    fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
         unsafe {
             let descr = PY_ARRAY_API.PyArray_DescrFromType(py, npy_type as _);
             py.from_owned_ptr(descr as _)
         }
     }
 
+    pub(crate) fn new_from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
+        unsafe {
+            let descr = PY_ARRAY_API.PyArray_DescrNewFromType(py, npy_type as _);
+            py.from_owned_ptr(descr as _)
+        }
+    }
+
     /// Returns the [array scalar][arrays-scalars] corresponding to this type descriptor.
     ///
     /// Equivalent to [`numpy.dtype.type`][dtype-type].

src/lib.rs

@@ -39,6 +39,7 @@
 pub mod array;
 pub mod borrow;
 pub mod convert;
+pub mod datetime;
 mod dtype;
 mod error;
 pub mod npyffi;