HDF5Exporter handling of ragged curves with tests (#1072)

* HDF5Exporter handles ragged curves by saving them into different datasets based on their names.

* Add HDF5Exporter tests

* Document HDF5Exporter

* Fix tmp file path

azure-test-template.yml

@@ -116,9 +116,9 @@ jobs:
           conda install -c conda-forge $(qt.bindings) --yes
         fi
         echo "Installing remainder of dependencies"
-        conda install -c conda-forge numpy scipy six pyopengl --yes
+        conda install -c conda-forge numpy scipy six pyopengl h5py --yes
       else
-        pip install $(qt.bindings) numpy scipy pyopengl six
+        pip install $(qt.bindings) numpy scipy pyopengl six h5py
       fi
       echo ""
       pip install pytest pytest-xdist pytest-cov coverage  

doc/source/exporting.rst

@@ -30,8 +30,13 @@ Export Formats
   for export.
 * Printer - Exports to the operating system's printing service. This exporter is provided for completeness, 
   but is not well supported due to problems with Qt's printing system.
+* HDF5 - Exports data from a :class:`~pyqtgraph.PlotItem` to a HDF5 file if
+  h5py_ is installed. This exporter supports :class:`~pyqtgraph.PlotItem`
+  objects containing multiple curves, stacking the data into a single HDF5
+  dataset based on the ``columnMode`` parameter. If data items aren't the same
+  size, each one is given its own dataset.
 
-
+.. _h5py: https://www.h5py.org/
 
 Exporting from the API
 ----------------------

pyqtgraph/exporters/HDF5Exporter.py

@@ -44,20 +44,27 @@ class HDF5Exporter(Exporter):
         data = []
 
         appendAllX = self.params['columnMode'] == '(x,y) per plot'
-        #print dir(self.item.curves[0])
+        # Check if the arrays are ragged
-        tlen = 0
+        len_first = len(self.item.curves[0].getData()[0]) if self.item.curves[0] else None
+        ragged = any(len(i.getData()[0]) != len_first for i in self.item.curves)
+
+        if ragged:
+            dgroup = fd.create_group(dsname)
+            for i, c in enumerate(self.item.curves):
+                d = c.getData()
+                fdata = numpy.array([d[0], d[1]]).astype('double')
+                cname = c.name() if c.name() is not None else str(i)
+                dset = dgroup.create_dataset(cname, data=fdata)
+        else:
             for i, c in enumerate(self.item.curves):
                 d = c.getData()
-            if i > 0 and len(d[0]) != tlen:
-                raise ValueError ("HDF5 Export requires all curves in plot to have same length")
                 if appendAllX or i == 0:
                     data.append(d[0])
-                tlen = len(d[0])
                 data.append(d[1])
 
-
             fdata = numpy.array(data).astype('double')
             dset = fd.create_dataset(dsname, data=fdata)
+
         fd.close()
 
 if HAVE_HDF5:

pyqtgraph/exporters/tests/test_hdf5.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+import pytest
+import pyqtgraph as pg
+from pyqtgraph.exporters import HDF5Exporter
+import numpy as np
+from numpy.testing import assert_equal
+import h5py
+import os
+
+
+@pytest.fixture
+def tmp_h5(tmp_path):
+    yield tmp_path / "data.h5"
+
+
+@pytest.mark.parametrize("combine", [False, True])
+def test_HDF5Exporter(tmp_h5, combine):
+    # Basic test of functionality: multiple curves with shared x array. Tests
+    # both options for stacking the data (columnMode).
+    x = np.linspace(0, 1, 100)
+    y1 = np.sin(x)
+    y2 = np.cos(x)
+
+    plt = pg.plot()
+    plt.plot(x=x, y=y1)
+    plt.plot(x=x, y=y2)
+
+    ex = HDF5Exporter(plt.plotItem)
+
+    if combine:
+        ex.parameters()['columnMode'] = '(x,y,y,y) for all plots'
+
+    ex.export(fileName=tmp_h5)
+
+    with h5py.File(tmp_h5, 'r') as f:
+        # should be a single dataset with the name of the exporter
+        dset = f[ex.parameters()['Name']]
+        assert isinstance(dset, h5py.Dataset)
+
+        if combine:
+            assert_equal(np.array([x, y1, y2]), dset)
+        else:
+            assert_equal(np.array([x, y1, x, y2]), dset)
+
+
+def test_HDF5Exporter_unequal_lengths(tmp_h5):
+    # Test export with multiple curves of different size. The exporter should
+    # detect this and create multiple hdf5 datasets under a group.
+    x1 = np.linspace(0, 1, 10)
+    y1 = np.sin(x1)
+    x2 = np.linspace(0, 1, 100)
+    y2 = np.cos(x2)
+
+    plt = pg.plot()
+    plt.plot(x=x1, y=y1, name='plot0')
+    plt.plot(x=x2, y=y2)
+
+    ex = HDF5Exporter(plt.plotItem)
+    ex.export(fileName=tmp_h5)
+
+    with h5py.File(tmp_h5, 'r') as f:
+        # should be a group with the name of the exporter
+        group = f[ex.parameters()['Name']]
+        assert isinstance(group, h5py.Group)
+
+        # should be a dataset under the group with the name of the PlotItem
+        assert_equal(np.array([x1, y1]), group['plot0'])
+
+        # should be a dataset under the group with a default name that's the
+        # index of the curve in the PlotItem
+        assert_equal(np.array([x2, y2]), group['1'])