starting the implementation of NexusEntry class

src/tavi/data/nxentry.py

@@ -0,0 +1,65 @@
+import h5py
+
+
+class NexusEntry(dict):
+
+    @staticmethod
+    def _getitem_recursively(obj, key, ATTRS):
+        "find key in obj recursively, return None if nonexsiting"
+        value = None
+        if key in obj:
+            if ATTRS:
+                try:
+                    value = obj[key]["attrs"]
+                except KeyError:
+                    print("Attribute does not exist.")
+            else:
+                try:
+                    value = obj[key]["dataset"]
+                except KeyError:
+                    print("Dataset does not exist.")
+
+        for k, v in obj.items():
+            if value is not None:
+                break
+            if k == "attrs" or k == "dataset":
+                continue
+
+            if isinstance(v, dict):
+                value = NexusEntry._getitem_recursively(v, key, ATTRS)
+
+        return value
+
+    def get(self, key, ATTRS=False, default=None):
+        """
+        Return dataset spicified by key regardless of the hierarchy.
+        Return attributes instead if ATTRS is True.
+
+        Note:
+            Only works if the key is unique
+        """
+        value = NexusEntry._getitem_recursively(self, key, ATTRS)
+        if value is not None:
+            return value
+        else:
+            return default
+
+    @staticmethod
+    def _write_recursively(items, nexus_entry):
+        """write items to nexus entry recursively"""
+        for key, value in items.items():
+            if key == "attrs":
+                for attr_key, attr_value in value.items():
+                    nexus_entry.attrs[attr_key] = attr_value
+            else:
+                if isinstance(value, dict):
+                    if "dataset" in value.keys():
+                        ds = nexus_entry.create_dataset(name=key, data=value["dataset"], maxshape=None)
+                        NexusEntry._write_recursively(value, ds)
+                    else:
+                        grp = nexus_entry.create_group(key + "/")
+                        NexusEntry._write_recursively(value, grp)
+
+    def to_nexus(self, path_to_nexus):
+        with h5py.File(path_to_nexus, "w") as nexus_file:
+            NexusEntry._write_recursively(self, nexus_file)

tests/test_nxentry.py

@@ -0,0 +1,121 @@
+import h5py
+import numpy as np
+import pytest
+
+from tavi.data.nxentry import NexusEntry
+
+
+def test_get_dataset(nexus_entry):
+    assert nexus_entry.get("definition") == "NXtas"
+    assert np.allclose(nexus_entry.get("a2"), np.array([242.0, 242.1, 242.2]))
+    assert nexus_entry.get("data", ATTRS=True) == {
+        "EX_required": "true",
+        "NX_class": "NXdata",
+        "axes": "en",
+        "signal": "detector",
+    }
+    assert nexus_entry.get("a3") is None
+    assert nexus_entry.get("detector") is None
+
+
+def test_to_nexus(nexus_entry):
+    path_to_nexus_entry = "./test_data/scan_to_nexus_test.h5"
+    nexus_entry.to_nexus(path_to_nexus_entry)
+
+    with h5py.File(path_to_nexus_entry, "r") as nexus_file:
+        assert str(nexus_file["scan0034"]["title"].asstr()[...]) == "scan_title"
+        assert nexus_file["scan0034"].attrs["NX_class"] == "NXentry"
+
+
+@pytest.fixture
+def nexus_entry():
+
+    analyser = {
+        "attrs": {
+            "EX_required": "true",
+            "NX_class": "NXcrystal",
+            "type": "NX_CHAR",
+        },
+        "a1": {
+            "attrs": {
+                "EX_required": "true",
+                "units": "degrees",
+            },
+            "dataset": np.array([142.0, 142.1, 142.2]),
+        },
+        "a2": {
+            "attrs": {
+                "EX_required": "true",
+                "units": "degrees",
+            },
+            "dataset": np.array([242.0, 242.1, 242.2]),
+        },
+        "sense": {"dataset": "-"},
+        "type": {"dataset": "Pg002"},
+    }
+    detector = {
+        "attrs": {
+            "EX_required": "true",
+            "NX_class": "NXdetector",
+        },
+        "data": {
+            "attrs": {
+                "EX_required": "true",
+                "type": "NX_INT",
+                "units": "counts",
+            },
+        },
+    }
+
+    instrument = {
+        "analyser": analyser,
+        "detector": detector,
+    }
+    monitor = {
+        "data": {
+            "attrs": {
+                "EX_required": "true",
+                "type": "NX_FLOAT",
+            },
+            "dataset": np.array([60, 60, 60]),
+        },
+        "attrs": {
+            "EX_required": "true",
+            "NX_class": "NXinstrument",
+        },
+    }
+
+    nexus_entry = NexusEntry(
+        [
+            (
+                "scan0034",
+                {
+                    "attrs": {
+                        "EX_required": "true",
+                        "NX_class": "NXentry",
+                        "dataset_name": "IPTS32124_CG4C_exp0424",
+                    },
+                    "definition": {
+                        "attrs": {"EX_required": "true", "type": "NX_CHAR"},
+                        "dataset": "NXtas",
+                    },
+                    "title": {
+                        "attrs": {"EX_required": "true", "type": "NX_CHAR"},
+                        "dataset": "scan_title",
+                    },
+                    "data": {
+                        "attrs": {
+                            "EX_required": "true",
+                            "NX_class": "NXdata",
+                            "axes": "en",
+                            "signal": "detector",
+                        }
+                    },
+                    "instrument": instrument,
+                    "monitor": monitor,
+                },
+            )
+        ]
+    )
+
+    return nexus_entry