Merge pull request #640 from European-XFEL/alias-links

Add support for generating Jupyterhub links to alias files

Author: JamesWrigley

docs/reading_files.rst

@@ -132,8 +132,10 @@ will print something like this (note the ✗ for an invalid alias)::
       agipd3-data: image.data
       agipd3-mask: image.mask
 
-Aliases can be used for selections as well, see
-:meth:`DataCollection.AliasIndexer.select` for more details.
+If aliases were loaded from a file, then clickable links to them will be
+displayed in the ``run.alias`` output (or you can call
+``run.alias.jhub_links()`` explicitly). Aliases can be used for selections as
+well, see :meth:`DataCollection.AliasIndexer.select` for more details.
 
 
 

extra_data/aliases.py

@@ -93,8 +93,21 @@ def source_alias(source):
         if len(source_key_aliases) == 0:
             return "No aliases have been loaded."
 
+        # Print links to the alias files
+        output_lines = []
+        n_files = len(self.file_paths())
+        links = self.jhub_links()
+        if n_files == 1:
+            path = self.file_paths()[0]
+            output_lines.append(f"Alias file: {links.get(path, path)}\n")
+        elif n_files > 1:
+            output_lines.append(f"Alias files:")
+            for i, path in enumerate(self.file_paths()):
+                line_end = "\n" if i == n_files - 1 else ""
+                output_lines.append(f"- {links.get(path, path)}{line_end}")
+
         # Print the aliases
-        output_lines = ["Loaded aliases:"]
+        output_lines.append("Loaded aliases:")
         for source, alias_keys in source_key_aliases.items():
             if len(alias_keys) == 0:
                 # If there are no keys then this is a plain source alias
@@ -167,6 +180,23 @@ def _resolve_aliased_selection(self, selection):
 
         return res
 
+    def file_paths(self):
+        """Return any file paths that were used to add aliases."""
+        return self.data._alias_files
+
+    def jhub_links(self):
+        """Return a dict of alias file paths to clickable Jupyterhub links.
+
+        Note that a link will only be generated if the file path is under ``/gpfs``.
+        """
+        links = { }
+        for p in self.file_paths():
+            path_from_home = str(p).replace("/gpfs", "GPFS")
+            if path_from_home.startswith("GPFS"):
+                links[p] = f"https://max-jhub.desy.de/hub/user-redirect/lab/tree/{path_from_home}"
+
+        return links
+
     def select(self, seln_or_alias, key_glob='*', require_all=False,
                require_any=False):
         """Select a subset of sources and keys from this data using aliases.

extra_data/reader.py

@@ -22,7 +22,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
-from itertools import groupby
+from itertools import groupby, chain
 from multiprocessing import Pool
 from operator import index
 from pathlib import Path
@@ -80,6 +80,7 @@ class DataCollection:
     def __init__(
             self, files, sources_data=None, train_ids=None, aliases=None,
             ctx_closes=False, *, inc_suspect_trains=True, is_single_run=False,
+            alias_files=None
     ):
         self.files = list(files)
         self.ctx_closes = ctx_closes
@@ -117,6 +118,9 @@ def __init__(
             }
         self._sources_data = sources_data
 
+        # Note that _alias_files is only for tracking where the aliases came
+        # from, the actual aliases are stored in _aliases.
+        self._alias_files = [] if alias_files is None else alias_files
         self._aliases = aliases or {}
         self.alias = AliasIndexer(self)
 
@@ -629,6 +633,10 @@ def _merge_aliases(self, alias_dicts):
 
         return new_aliases
 
+    def _merge_alias_files(self, *alias_files):
+        all_files = chain.from_iterable(alias_files)
+        return sorted(set(all_files))
+
     def union(self, *others):
         """Join the data in this collection with one or more others.
 
@@ -653,6 +661,8 @@ def union(self, *others):
 
         aliases = self._merge_aliases(
             [self._aliases] + [dc._aliases for dc in others])
+        alias_files = self._merge_alias_files(self._alias_files,
+                                              *[dc._alias_files for dc in others])
 
         train_ids = sorted(set().union(*[sd.train_ids for sd in sources_data.values()]))
         # Update the internal list of train IDs for the sources
@@ -664,7 +674,7 @@ def union(self, *others):
         return DataCollection(
             files, sources_data=sources_data, train_ids=train_ids,
             aliases=aliases, inc_suspect_trains=self.inc_suspect_trains,
-            is_single_run=same_run(self, *others),
+            is_single_run=same_run(self, *others), alias_files=alias_files
         )
 
     def __or__(self, other):
@@ -677,6 +687,7 @@ def _parse_aliases(self, alias_defs):
         """Parse alias definitions into alias dictionaries."""
 
         alias_dicts = []
+        alias_files = []
 
         def is_valid_alias(k, v):
             return (isinstance(k, str) and (
@@ -693,10 +704,11 @@ def is_valid_alias(k, v):
                 alias_dicts.append(alias_def)
             elif isinstance(alias_def, (str, os.PathLike)):
                 # From a file.
+                alias_files.append(Path(alias_def))
                 alias_dicts.append(
                     self._load_aliases_from_file(Path(alias_def)))
 
-        return alias_dicts
+        return alias_dicts, alias_files
 
     def _load_aliases_from_file(self, aliases_path):
         """Load alias definitions from file."""
@@ -786,14 +798,15 @@ def with_aliases(self, *alias_defs):
         """
 
         # Check for conflicts within these definitions
-        new_aliases = self._merge_aliases(
-            [self._aliases] + self._parse_aliases(alias_defs))
+        new_aliases, new_alias_files = self._parse_aliases(alias_defs)
+        new_aliases = self._merge_aliases([self._aliases] + new_aliases)
+        alias_files = self._merge_alias_files(self._alias_files, new_alias_files)
 
         return DataCollection(
             self.files, sources_data=self._sources_data,
             train_ids=self.train_ids, aliases=new_aliases,
             inc_suspect_trains=self.inc_suspect_trains,
-            is_single_run=self.is_single_run
+            is_single_run=self.is_single_run, alias_files=alias_files
         )
 
     def only_aliases(self, *alias_defs, strict=False, require_all=False):
@@ -818,8 +831,9 @@ def only_aliases(self, *alias_defs, strict=False, require_all=False):
         """
 
         # Create new aliases.
-        aliases = self._merge_aliases(
-            [self._aliases] + self._parse_aliases(alias_defs))
+        new_aliases, new_alias_files = self._parse_aliases(alias_defs)
+        aliases = self._merge_aliases([self._aliases] + new_aliases)
+        alias_files = self._merge_alias_files(self._alias_files, new_alias_files)
 
         # Set of sources aliased.
         aliased_sources = {literal for literal in aliases.values()
@@ -865,6 +879,7 @@ def only_aliases(self, *alias_defs, strict=False, require_all=False):
         # Create a new DataCollection from selecting and add the aliases.
         new_data = self.select(selection, require_all=require_all)
         new_data._aliases = aliases
+        new_data._alias_files = alias_files
 
         return new_data
 
@@ -1092,7 +1107,7 @@ def select(self, seln_or_source_glob, key_glob='*', require_all=False,
         return DataCollection(
             files, sources_data, train_ids=train_ids, aliases=self._aliases,
             inc_suspect_trains=self.inc_suspect_trains,
-            is_single_run=self.is_single_run
+            is_single_run=self.is_single_run, alias_files=self._alias_files
         )
 
     def deselect(self, seln_or_source_glob, key_glob='*'):
@@ -1129,7 +1144,7 @@ def deselect(self, seln_or_source_glob, key_glob='*'):
         return DataCollection(
             files, sources_data=sources_data, train_ids=self.train_ids,
             aliases=self._aliases, inc_suspect_trains=self.inc_suspect_trains,
-            is_single_run=self.is_single_run,
+            is_single_run=self.is_single_run, alias_files=self._alias_files
         )
 
     def select_trains(self, train_range):
@@ -1164,7 +1179,7 @@ def select_trains(self, train_range):
         return DataCollection(
             files, sources_data=sources_data, train_ids=new_train_ids,
             aliases=self._aliases, inc_suspect_trains=self.inc_suspect_trains,
-            is_single_run=self.is_single_run,
+            is_single_run=self.is_single_run, alias_files=self._alias_files
         )
 
     def split_trains(self, parts=None, trains_per_part=None):
@@ -1207,7 +1222,7 @@ def dict_zip(iter_d):
             yield DataCollection(
                 files, sources_data=sources_data_part, train_ids=train_ids,
                 aliases=self._aliases, inc_suspect_trains=self.inc_suspect_trains,
-                is_single_run=self.is_single_run,
+                is_single_run=self.is_single_run, alias_files=self._alias_files
             )
 
     def _check_source_conflicts(self):

extra_data/tests/conftest.py

@@ -5,10 +5,22 @@
 import numpy as np
 import pytest
 from tempfile import TemporaryDirectory
+import yaml
 
 from . import make_examples
 
 
+# Make Pyyaml write tuples as regular lists, otherwise it adds a tag to the
+# element that makes it impossible to deserialize without a custom constructor.
+def tuple_representer(dumper, data):
+    return dumper.represent_list(list(data))
+
+class SafeDumperWithTuples(yaml.SafeDumper):
+     pass
+
+yaml.add_representer(tuple, tuple_representer, Dumper=SafeDumperWithTuples)
+
+
 @pytest.fixture(scope='session', params=['0.5', '1.0', '1.2'])
 def format_version(request):
     return request.param
@@ -78,6 +90,14 @@ def mock_sa3_control_aliases():
         'bogus-key': ('SA3_XTD10_XGM/XGM/DOOCS', 'foo')
     }
 
+@pytest.fixture
+def mock_sa3_control_aliases_yaml(mock_sa3_control_aliases, tmp_path):
+    aliases_path = tmp_path / "mock_sa3_control_aliases.yaml"
+    with aliases_path.open("w") as f:
+        yaml.dump(mock_sa3_control_aliases, f, Dumper=SafeDumperWithTuples)
+
+    return aliases_path
+
 
 @pytest.fixture(scope='module')
 def mock_control_data_with_empty_source(format_version):