fix: writer support for numpy scalars (#72)

Author: terraputix

src/reader.rs

@@ -6,8 +6,8 @@ use crate::{
 use delegate::delegate;
 use num_traits::Zero;
 use numpy::{
-    ndarray::{self, Array0},
-    Element, PyArray0,
+    ndarray::{self},
+    Element,
 };
 use omfiles_rs::{
     reader::{OmFileArray as OmFileArrayRs, OmFileReader as OmFileReaderRs},
@@ -124,10 +124,23 @@ impl OmFileReader {
                 .map_err(|_| Self::only_scalars_error())?;
 
             let value = scalar_reader.read_scalar::<T>();
-            let array_base = Array0::from_elem([], value.unwrap());
-            let py_scalar = PyArray0::from_owned_array(py, array_base);
 
-            return Ok(py_scalar.into_any().unbind());
+            let numpy = py.import("numpy")?;
+            let np_type = match std::any::type_name::<T>() {
+                "f32" => numpy.getattr("float32")?,
+                "f64" => numpy.getattr("float64")?,
+                "i8" => numpy.getattr("int8")?,
+                "u8" => numpy.getattr("uint8")?,
+                "i16" => numpy.getattr("int16")?,
+                "u16" => numpy.getattr("uint16")?,
+                "i32" => numpy.getattr("int32")?,
+                "u32" => numpy.getattr("uint32")?,
+                "i64" => numpy.getattr("int64")?,
+                "u64" => numpy.getattr("uint64")?,
+                _ => return Err(PyErr::new::<PyValueError, _>("Unsupported type")),
+            };
+            let py_scalar = np_type.call1((value,))?;
+            Ok(py_scalar.into())
         })
     }
 }
@@ -536,7 +549,7 @@ impl OmFileReader {
     /// Read the scalar value of the variable.
     ///
     /// Returns:
-    ///     object: The scalar value as a Python object (str, int, or float).
+    ///     object: The scalar value as a numpy scalar or a Python string.
     ///
     /// Raises:
     ///     ValueError: If the variable is not a scalar.

src/reader_async.rs

@@ -7,8 +7,8 @@ use async_lock::RwLock;
 use delegate::delegate;
 use num_traits::Zero;
 use numpy::{
-    ndarray::{self, Array0},
-    Element, PyArray0,
+    ndarray::{self},
+    Element,
 };
 use omfiles_rs::{
     reader_async::OmFileReaderAsync as OmFileReaderAsyncRs,
@@ -135,10 +135,23 @@ impl OmFileReaderAsync {
                 .map_err(|_| Self::only_scalars_error())?;
 
             let value = scalar_reader.read_scalar::<T>();
-            let array_base = Array0::from_elem([], value.unwrap());
-            let py_scalar = PyArray0::from_owned_array(py, array_base);
 
-            return Ok(py_scalar.into_any().unbind());
+            let numpy = py.import("numpy")?;
+            let np_type = match std::any::type_name::<T>() {
+                "f32" => numpy.getattr("float32")?,
+                "f64" => numpy.getattr("float64")?,
+                "i8" => numpy.getattr("int8")?,
+                "u8" => numpy.getattr("uint8")?,
+                "i16" => numpy.getattr("int16")?,
+                "u16" => numpy.getattr("uint16")?,
+                "i32" => numpy.getattr("int32")?,
+                "u32" => numpy.getattr("uint32")?,
+                "i64" => numpy.getattr("int64")?,
+                "u64" => numpy.getattr("uint64")?,
+                _ => return Err(PyErr::new::<PyValueError, _>("Unsupported type")),
+            };
+            let py_scalar = np_type.call1((value,))?;
+            Ok(py_scalar.into())
         })
     }
 }
@@ -486,7 +499,7 @@ impl OmFileReaderAsync {
     /// Read the scalar value of the variable.
     ///
     /// Returns:
-    ///     object: The scalar value as a Python object (str, int, or float).
+    ///     object: The scalar value as a numpy scalar or a Python string.
     ///
     /// Raises:
     ///     ValueError: If the variable is not a scalar.

src/writer.rs

@@ -332,6 +332,35 @@ impl OmFileWriter {
             .map(Into::into)
             .collect();
 
+        let py = value.py();
+
+        // make an instance check against numpy scalar types
+        macro_rules! check_numpy_type {
+            ($numpy:expr, $type_name:literal, $rust_type:ty) => {
+                if let Ok(numpy_type) = $numpy.getattr($type_name) {
+                    if value.is_instance(&numpy_type)? {
+                        let scalar_value: $rust_type = value.call_method0("item")?.extract()?;
+                        return self.store_scalar(scalar_value, name, &children);
+                    }
+                }
+            };
+        }
+
+        // Try to import numpy and check for numpy scalar types
+        if let Ok(numpy) = py.import("numpy") {
+            check_numpy_type!(numpy, "int8", i8);
+            check_numpy_type!(numpy, "uint8", u8);
+            check_numpy_type!(numpy, "int16", i16);
+            check_numpy_type!(numpy, "uint16", u16);
+            check_numpy_type!(numpy, "int32", i32);
+            check_numpy_type!(numpy, "uint32", u32);
+            check_numpy_type!(numpy, "int64", i64);
+            check_numpy_type!(numpy, "uint64", u64);
+            check_numpy_type!(numpy, "float32", f32);
+            check_numpy_type!(numpy, "float64", f64);
+        }
+
+        // Fall back to Python built-in types
         let result = if let Ok(_value) = value.extract::<String>() {
             self.store_scalar(value.to_string(), name, &children)?
         } else if let Ok(value) = value.extract::<f64>() {

tests/test_read_write.py

@@ -71,8 +71,8 @@ def test_write_hierarchical_file(empty_temp_om_file):
     child2_var = writer.write_array(child2_data, chunks=[1, 1], name="child2", scale_factor=100000.0)
 
     # Write attributes and get their variables
-    meta1_var = writer.write_scalar(42.0, name="metadata1")
-    meta2_var = writer.write_scalar(123, name="metadata2")
+    meta1_var = writer.write_scalar(np.float32(42.0), name="metadata1")
+    meta2_var = writer.write_scalar(np.int32(123), name="metadata2")
     meta3_var = writer.write_scalar("blub", name="metadata3")
 
     # Write child1 array with attribute children
@@ -115,21 +115,30 @@ def test_write_hierarchical_file(empty_temp_om_file):
     assert read_child2.dtype == np.float32
 
     # Verify metadata attributes
-    metadata_reader0 = child1_reader.get_child_by_index(0)
-    metadata = metadata_reader0.read_scalar()
-    assert metadata == 42.0
-    assert metadata_reader0.dtype == np.float64
-
-    metadata_reader2 = child1_reader.get_child_by_index(2)
-    metadata = metadata_reader2.read_scalar()
-    assert metadata == "blub"
-    assert metadata_reader2.dtype == str
+    metadata_reader1 = child1_reader.get_child_by_index(0)
+    metadata1 = metadata_reader1.read_scalar()
+    assert metadata1 == 42.0
+    assert type(metadata1) == np.float32
+    assert metadata_reader1.dtype == np.float32
+
+    metadata_reader2 = child1_reader.get_child_by_index(1)
+    metadata2 = metadata_reader2.read_scalar()
+    assert metadata2 == 123
+    assert type(metadata2) == np.int32
+    assert metadata_reader2.dtype == np.int32
+
+    metadata_reader3 = child1_reader.get_child_by_index(2)
+    metadata3 = metadata_reader3.read_scalar()
+    assert metadata3 == "blub"
+    assert type(metadata3) == str
+    assert metadata_reader3.dtype == str
 
     reader.close()
     child1_reader.close()
     child2_reader.close()
-    metadata_reader0.close()
+    metadata_reader1.close()
     metadata_reader2.close()
+    metadata_reader3.close()
 
 
 @pytest.mark.asyncio