Add tests for aliases

extra_data/tests/conftest.py

@@ -56,6 +56,29 @@ def mock_sa3_control_data(format_version):
         yield path
 
 
+@pytest.fixture
+def mock_sa3_control_aliases():
+    return {
+        'sa3-xgm': 'SA3_XTD10_XGM/XGM/DOOCS',
+        'hv': ('SA3_XTD10_XGM/XGM/DOOCS', 'pulseEnergy.wavelengthUsed'),
+        'beam-x': ('SA3_XTD10_XGM/XGM/DOOCS', 'beamPosition.ixPos'),
+        'beam-y': ('SA3_XTD10_XGM/XGM/DOOCS', 'beamPosition.iyPos'),
+
+        'imgfel-frames': ('SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput', 'data.image.pixels'),
+        'imgfel-frames2': ('SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput', 'data.image.pixels'),
+        'imgfel-screen-pos': ('SA3_XTD10_IMGFEL/MOTOR/SCREEN', 'actualPosition'),
+        'imgfel-filter-pos': ('SA3_XTD10_IMGFEL/MOTOR/FILTER', 'actualPosition'),
+
+        'mcp-adc': 'SA3_XTD10_MCP/ADC/1',
+        'mcp-mpod': 'SA3_XTD10_MCP/MCPS/MPOD',
+        'mcp-voltage': ('SA3_XTD10_MCP/MCPS/MPOD', 'channels.U3.voltage'),
+        'mcp-trace': ('SA3_XTD10_MCP/ADC/1:channel_5.output', 'data.rawData'),
+
+        'bogus-source': 'SA4_XTD20_XGM/XGM/DOOCS',
+        'bogus-key': ('SA3_XTD10_XGM/XGM/DOOCS', 'foo')
+    }
+
+
 @pytest.fixture(scope='module')
 def mock_control_data_with_empty_source(format_version):
     with TemporaryDirectory() as td:

extra_data/tests/test_reader_mockdata.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta, timezone
 from itertools import islice
 from multiprocessing import Process
+from pathlib import Path
 from warnings import catch_warnings
 
 import h5py
@@ -10,15 +11,15 @@
 import pytest
 import stat
 import shutil
-from tempfile import mkdtemp
+from tempfile import mkdtemp, TemporaryDirectory
 from testpath import assert_isfile
 from unittest import mock
 from xarray import DataArray
 
 from extra_data import (
     H5File, RunDirectory, by_index, by_id,
-    SourceNameError, PropertyNameError, DataCollection, open_run,
-    MultiRunError
+    DataCollection, KeyData, open_run,
+    SourceNameError, PropertyNameError, AliasError, MultiRunError
 )
 
 def test_iterate_trains(mock_agipd_data, mock_control_data_with_empty_source):
@@ -990,3 +991,273 @@ def test_run_metadata_no_trains(mock_scs_run):
     sel = run.select_trains(np.s_[:0])
     md = sel.run_metadata()
     assert md['dataFormatVersion'] == '1.0'
+
+
+def test_with_aliases(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+
+    def assert_equal_keydata(kd1, kd2):
+        assert isinstance(kd1, KeyData)
+        assert isinstance(kd2, KeyData)
+        assert kd1.source == kd2.source
+        assert kd1.key == kd2.key
+        assert kd1.train_ids == kd2.train_ids
+
+    # Test whether source alias yields identical SourceData.
+    assert run.alias['sa3-xgm'] is run['SA3_XTD10_XGM/XGM/DOOCS']
+
+    # Test whether source alias plus literal key yields equal KeyData.
+    assert_equal_keydata(
+        run.alias['sa3-xgm', 'pulseEnergy.wavelengthUsed'],
+        run['SA3_XTD10_XGM/XGM/DOOCS', 'pulseEnergy.wavelengthUsed'])
+
+    # Test whether key alias yields equal KeyData.
+    assert_equal_keydata(
+        run.alias['hv'],
+        run['SA3_XTD10_XGM/XGM/DOOCS', 'pulseEnergy.wavelengthUsed'])
+
+    # Test undefined aliases.
+    with pytest.raises(AliasError):
+        run.alias['foo']
+        run.alias['foo', 'bar']
+
+    # Test using a literal key with a key alias.
+    with pytest.raises(ValueError):
+        run.alias['hv', 'pressure']
+
+    # Test using an existing source alias for a non-existing source.
+    with pytest.raises(SourceNameError):
+        run.alias['bogus-source']
+
+    # Test using an existing key alias for a non-existing key.
+    with pytest.raises(PropertyNameError):
+        run.alias['bogus-key']
+
+
+def test_json_alias_file(mock_sa3_control_data, mock_sa3_control_aliases):
+    with TemporaryDirectory() as td:
+        aliases_path = Path(td) / 'aliases.json'
+        aliases_path.write_text('''
+{
+    "sa3-xgm": "SA3_XTD10_XGM/XGM/DOOCS",
+    "SA3_XTD10_XGM/XGM/DOOCS": {
+        "hv": "pulseEnergy.wavelengthUsed",
+        "beam-x": "beamPosition.ixPos",
+        "beam-y": "beamPosition.iyPos"
+    },
+
+    "imgfel-frames": ["SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput", "data.image.pixels"],
+    "imgfel-frames2": ["SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput", "data.image.pixels"],
+    "imgfel-screen-pos": ["SA3_XTD10_IMGFEL/MOTOR/SCREEN", "actualPosition"],
+    "imgfel-filter-pos": ["SA3_XTD10_IMGFEL/MOTOR/FILTER", "actualPosition"],
+
+    "mcp-adc": "SA3_XTD10_MCP/ADC/1",
+    "mcp-mpod": "SA3_XTD10_MCP/MCPS/MPOD",
+    "mcp-voltage": ["SA3_XTD10_MCP/MCPS/MPOD", "channels.U3.voltage"],
+    "mcp-trace": ["SA3_XTD10_MCP/ADC/1:channel_5.output", "data.rawData"],
+
+    "bogus-source": "SA4_XTD20_XGM/DOOCS/MAIN",
+    "bogus-key": ["SA1_XTD2_XGM/DOOCS/MAIN", "foo"]
+}
+        ''')
+
+        run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+        assert run._aliases == mock_sa3_control_aliases
+
+
+def test_yaml_alias_file(mock_sa3_control_data, mock_sa3_control_aliases):
+    with TemporaryDirectory() as td:
+        aliases_path = Path(td) / 'aliases.yaml'
+        aliases_path.write_text('''
+sa3-xgm: SA3_XTD10_XGM/XGM/DOOCS
+
+SA3_XTD10_XGM/XGM/DOOCS:
+  hv: pulseEnergy.wavelengthUsed
+  beam-x: beamPosition.ixPos
+  beam-y: beamPosition.iyPos
+
+imgfel-frames: [SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput, data.image.pixels]
+imgfel-frames2: [SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput, data.image.pixels]
+imgfel-screen-pos: [SA3_XTD10_IMGFEL/MOTOR/SCREEN, actualPosition]
+imgfel-filter-pos: [SA3_XTD10_IMGFEL/MOTOR/FILTER, actualPosition]
+
+mcp-adc: SA3_XTD10_MCP/ADC/1
+mcp-mpod: SA3_XTD10_MCP/MCPS/MPOD
+mcp-voltage: [SA3_XTD10_MCP/MCPS/MPOD", channels.U3.voltage]
+mcp-trace: [SA3_XTD10_MCP/ADC/1:channel_5.output, data.rawData]
+
+bogus-source: SA4_XTD20_XGM/DOOCS/MAIN
+bogus-key: [SA1_XTD2_XGM/DOOCS/MAIN, foo]
+        ''')
+
+        run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+        assert run._aliases == mock_sa3_control_aliases
+
+
+def test_toml_alias_file(mock_sa3_control_data, mock_sa3_control_aliases):
+    with TemporaryDirectory() as td:
+        aliases_path = Path(td) / 'aliases.toml'
+        aliases_path.write_text('''
+sa3-xgm = "SA3_XTD10_XGM/XGM/DOOCS"
+
+imgfel-frames = ["SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput", "data.image.pixels"]
+imgfel-frames2 = ["SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput", "data.image.pixels"]
+imgfel-screen-pos = ["SA3_XTD10_IMGFEL/MOTOR/SCREEN", "actualPosition"]
+imgfel-filter-pos = ["SA3_XTD10_IMGFEL/MOTOR/FILTER", "actualPosition"]
+
+mcp-adc = "SA3_XTD10_MCP/ADC/1"
+mcp-mpod = "SA3_XTD10_MCP/MCPS/MPOD"
+mcp-voltage = ["SA3_XTD10_MCP/MCPS/MPOD", "channels.U3.voltage"]
+mcp-trace = ["SA3_XTD10_MCP/ADC/1:channel_5.output", "data.rawData"]
+
+bogus-source = "SA4_XTD20_XGM/DOOCS/MAIN"
+bogus-key = ["SA1_XTD2_XGM/DOOCS/MAIN", "foo"]
+
+["SA3_XTD10_XGM/XGM/DOOCS"]
+hv = "pulseEnergy.wavelengthUsed"
+beam-x = "beamPosition.ixPos"
+beam-y = "beamPosition.iyPos"
+        ''')
+
+        run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+        assert run._aliases == mock_sa3_control_aliases
+
+
+def test_only_aliases(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+    subrun = H5File(mock_sa3_control_data).only_aliases(mock_sa3_control_aliases)
+
+    # Assume that aliases work when the _aliases property is equal.
+    assert run._aliases == subrun._aliases
+
+    # Test whether only the sources used in aliases are present.
+    assert subrun.all_sources == {
+        'SA3_XTD10_XGM/XGM/DOOCS',
+        'SA3_XTD10_IMGFEL/CAM/BEAMVIEW:daqOutput',
+        'SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput',
+        'SA3_XTD10_IMGFEL/MOTOR/SCREEN',
+        'SA3_XTD10_IMGFEL/MOTOR/FILTER',
+        'SA3_XTD10_MCP/ADC/1',
+        'SA3_XTD10_MCP/MCPS/MPOD',
+        'SA3_XTD10_MCP/ADC/1:channel_5.output',
+    }
+
+    # Test whether all keys are present for an aliased source.
+    assert subrun['SA3_XTD10_XGM/XGM/DOOCS'].keys() == run['SA3_XTD10_XGM/XGM/DOOCS'].keys()
+
+    # Test whether all keys are present for an aliased source, even if
+    # there are key aliases for it as well.
+    assert subrun['SA3_XTD10_MCP/MCPS/MPOD'].keys() == run['SA3_XTD10_MCP/MCPS/MPOD'].keys()
+
+    # Test whether only aliased keys are present for unaliased sources.
+    assert subrun['SA3_XTD10_IMGFEL/MOTOR/SCREEN'].keys() == {'actualPosition.value'}
+
+    # Test strict selection.
+    with pytest.raises(ValueError):
+        H5File(mock_sa3_control_data).only_aliases(
+            mock_sa3_control_aliases, strict=True)
+
+    # Remove bogus aliases and test strict selection again.
+    strict_aliases = mock_sa3_control_aliases.copy()
+    del strict_aliases['bogus-source']
+    del strict_aliases['bogus-key']
+    H5File(mock_sa3_control_data).only_aliases(strict_aliases, strict=True)
+
+    # Prepare a run with less trains for a single source
+    # (SA3_XTD10_IMGFEL/CAM/BEAMVIEW2:daqOutput) by removing all sources
+    # without any trains.
+    run = H5File(mock_sa3_control_data) \
+        .deselect([('SA3_XTD10_MCP/ADC/1:*', '*'),
+                   ('SA3_XTD10_IMGFEL/CAM/BEAMVIEW:*', '*')])
+    del strict_aliases['mcp-trace']
+    del strict_aliases['imgfel-frames']
+
+    # Without strict alias selection and a bogus alias.
+    subrun = run.only_aliases(mock_sa3_control_aliases,
+                              require_all=True, strict=False)
+    np.testing.assert_array_equal(subrun.train_ids, run.train_ids[1::2])
+
+    # With strict alias selection.
+    subrun = run.only_aliases(strict_aliases, require_all=True, strict=True)
+    np.testing.assert_array_equal(subrun.train_ids, run.train_ids[1::2])
+
+
+def test_preserve_aliases(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+
+    # Test whether selection operations preserve aliases.
+    assert run.select_trains(by_index[:5])._aliases == run._aliases
+    assert run.select_trains(by_id[run.train_ids[:5]])._aliases == run._aliases
+    assert run.select('*')._aliases == run._aliases
+    assert run.deselect('*XGM*')._aliases == run._aliases
+    assert all([subrun._aliases == run._aliases
+                for subrun in run.split_trains(parts=5)])
+
+
+def test_aliases_union(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+
+    # Split the aliases into two halves and test the union.
+    run1 = run.with_aliases(dict(islice(mock_sa3_control_aliases.items(), 0, None, 2)))
+    run2 = run.with_aliases(dict(islice(mock_sa3_control_aliases.items(), 1, None, 2)))
+    assert run1.union(run2)._aliases == mock_sa3_control_aliases
+
+    # Split the run into two.
+    even_run = run.select_trains(by_id[run.train_ids[0::2]])
+    odd_run = run.select_trains(by_id[run.train_ids[1::2]])
+
+    # Test overlapping aliases with no conflict.
+    even_run.union(odd_run)
+
+    # Test conflicting aliases.
+    conflicting_aliases = mock_sa3_control_aliases.copy()
+    conflicting_aliases['hv'] = ('SA3_XTD10_XGM/XGM/DOOCS', 'pressure.pressure1')
+    with pytest.raises(ValueError):
+        even_run.union(odd_run.with_aliases(conflicting_aliases))
+
+
+def test_alias_select(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+
+    # Only source alias.
+    subrun = run.alias.select('sa3-xgm')
+    assert subrun.all_sources == {'SA3_XTD10_XGM/XGM/DOOCS'}
+    assert subrun.alias['sa3-xgm'].keys() == run.alias['sa3-xgm'].keys()
+
+    # Source alias and key glob.
+    subrun = run.alias.select('sa3-xgm', 'pressure.pressure*.value')
+    assert subrun.all_sources == {'SA3_XTD10_XGM/XGM/DOOCS'}
+    assert subrun.alias['sa3-xgm'].keys() == {
+        'pressure.pressure1.value', 'pressure.pressureFiltered.value'}
+
+    # Iterable of aliases and/or with key globs.
+    subrun = run.alias.select([('sa3-xgm', 'pressure.pressure*.value'),
+                               'beam-x', 'mcp-voltage'])
+    assert subrun.all_sources == {'SA3_XTD10_XGM/XGM/DOOCS', 'SA3_XTD10_MCP/MCPS/MPOD'}
+    assert subrun.alias['sa3-xgm'].keys() == {
+        'pressure.pressure1.value', 'pressure.pressureFiltered.value',
+        'beamPosition.ixPos.value'}
+    assert subrun.alias['mcp-mpod'].keys() == {'channels.U3.voltage.value'}
+
+    # Dictionary
+    subrun = run.alias.select({'sa3-xgm': None, 'mcp-mpod': {'channels.U1.voltage'}})
+    assert subrun.all_sources == {'SA3_XTD10_XGM/XGM/DOOCS', 'SA3_XTD10_MCP/MCPS/MPOD'}
+    assert subrun.alias['sa3-xgm'].keys() == run.alias['sa3-xgm'].keys()
+    assert subrun.alias['mcp-mpod'].keys() == {'channels.U1.voltage.value'}
+
+
+def test_alias_deselect(mock_sa3_control_data, mock_sa3_control_aliases):
+    run = H5File(mock_sa3_control_data).with_aliases(mock_sa3_control_aliases)
+
+    # De-select via alias.
+    subrun = run.alias.deselect([
+        ('sa3-xgm', 'pressure.*'), ('sa3-xgm', 'current.*'),
+        ('sa3-xgm', 'gasDosing.*'), ('sa3-xgm', 'gasSupply.*'),
+        ('sa3-xgm', 'pressure.*'), ('sa3-xgm', 'pulseEnergy.*'),
+        ('sa3-xgm', 'signalAdaption.*')
+    ])
+    assert subrun.all_sources == run.all_sources
+    assert subrun.alias['sa3-xgm'].keys() == {
+        'beamPosition.ixPos.value', 'beamPosition.ixPos.timestamp',
+        'beamPosition.iyPos.value', 'beamPosition.iyPos.timestamp',
+        'pollingInterval.value', 'pollingInterval.timestamp'}

setup.py

@@ -76,6 +76,8 @@ def find_version(*parts):
               'pytest',
               'pytest-cov',
               'testpath',
+              'pyyaml',
+              'tomli'
           ]
       },
       python_requires='>=3.6',