Merge pull request #3683 from t20100/pint-friendly

silx.io.dictdump: Added support of `pint` in `dicttoh5` and `dicttonx`

src/silx/io/dictdump.py

@@ -30,8 +30,11 @@
 import logging
 import numpy
 import os.path
-import sys
 import h5py
+try:
+    import pint
+except ImportError:
+    pint = None
 
 from .configdict import ConfigDict
 from .utils import is_group
@@ -64,6 +67,8 @@ def _prepare_hdf5_write_value(array_like):
         ``numpy.array()`` (`str`, `list`, `numpy.ndarray`…)
     :return: ``numpy.ndarray`` ready to be written as an HDF5 dataset
     """
+    if pint is not None and isinstance(array_like, pint.quantity.Quantity):
+        return numpy.array(array_like.magnitude)
     array = numpy.asarray(array_like)
     if numpy.issubdtype(array.dtype, numpy.bytes_):
         return numpy.array(array_like, dtype=vlen_bytes)
@@ -455,6 +460,7 @@ def nexus_to_h5_dict(
                     value = h5py.SoftLink(first)
             elif is_link(value):
                 key = key[1:]
+
         if isinstance(value, Mapping):
             # HDF5 group
             key_has_nx_class = add_nx_class and _has_nx_class(treedict, key)
@@ -463,9 +469,15 @@ def nexus_to_h5_dict(
                 parents=parents+(key,),
                 add_nx_class=add_nx_class,
                 has_nx_class=key_has_nx_class)
+
+        elif pint is not None and isinstance(value, pint.quantity.Quantity):
+            copy[key] = value.magnitude
+            copy[(key, "units")] = f"{value.units:~C}"
+
         else:
             # HDF5 dataset or link
             copy[key] = value
+
     if add_nx_class and not has_nx_class:
         _ensure_nx_class(copy, parents)
     return copy

src/silx/io/test/test_dictdump.py

@@ -1,5 +1,5 @@
 # /*##########################################################################
-# Copyright (C) 2016-2021 European Synchrotron Radiation Facility
+# Copyright (C) 2016-2022 European Synchrotron Radiation Facility
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -26,15 +26,21 @@
 __license__ = "MIT"
 __date__ = "17/01/2018"
 
-from collections import OrderedDict
-import numpy
+
+from collections import defaultdict, OrderedDict
+from copy import deepcopy
+from io import BytesIO
 import os
 import tempfile
 import unittest
-import h5py
-from copy import deepcopy
 
-from collections import defaultdict
+import h5py
+import numpy
+try:
+    import pint
+except ImportError:
+    pint = None
+import pytest
 
 from silx.utils.testutils import LoggingValidator
 
@@ -47,6 +53,13 @@
 from ..utils import h5py_read_dataset
 
 
+@pytest.fixture
+def tmp_h5py_file():
+    with BytesIO() as buffer:
+        with h5py.File(buffer, mode="w") as h5file:
+            yield h5file
+
+
 def tree():
     """Tree data structure as a recursive nested dictionary"""
     return defaultdict(tree)
@@ -511,6 +524,22 @@ def assert_append(update_mode):
         assert_append("replace")
 
 
+@pytest.mark.skipif(pint is None, reason="Require pint")
+def test_dicttoh5_pint(tmp_h5py_file):
+    ureg = pint.UnitRegistry()
+    treedict = {
+        "array_mm": pint.Quantity([1, 2, 3], ureg.mm),
+        "value_kg": 3 * ureg.kg,
+    }
+
+    dicttoh5(treedict, tmp_h5py_file)
+
+    result = h5todict(tmp_h5py_file)
+    assert set(treedict.keys()) == set(result.keys())
+    for key, value in treedict.items():
+        assert numpy.array_equal(result[key], value.magnitude)
+
+
 class TestH5ToDict(H5DictTestCase):
     def setUp(self):
         self.tempdir = tempfile.mkdtemp()
@@ -799,6 +828,22 @@ def assert_append(update_mode, add_nx_class=None):
         assert_append("replace", add_nx_class=True)
 
 
+@pytest.mark.skipif(pint is None, reason="Require pint")
+def test_dicttonx_pint(tmp_h5py_file):
+    ureg = pint.UnitRegistry()
+    treedict = {
+        "array_mm": pint.Quantity([1, 2, 3], ureg.mm),
+        "value_kg": 3 * ureg.kg,
+    }
+
+    dictdump.dicttonx(treedict, tmp_h5py_file)
+
+    result = dictdump.nxtodict(tmp_h5py_file)
+    for key, value in treedict.items():
+        assert numpy.array_equal(result[key], value.magnitude)
+        assert result[f"{key}@units"] == f"{value.units:~C}"
+
+
 class TestNxToDict(H5DictTestCase):
     def setUp(self):
         self.tempdir = tempfile.mkdtemp()