from collections import defaultdict from collections.abc import Iterable from .exceptions import AliasError class AliasIndexer: __slots__ = ('data',) def __init__(self, data): self.data = data def _resolve_any_alias(self, alias): alias = alias.lower().replace('_', '-') try: literal = self.data._aliases[alias] except KeyError: raise AliasError(alias) from None return literal def _resolve_source_alias(self, alias): source = self._resolve_any_alias(alias) if isinstance(source, tuple): raise ValueError(f'{alias} not aliasing a source for this data') return source def __getitem__(self, aliased_item): if isinstance(aliased_item, tuple) and len(aliased_item) == 2: # Source alias with key literal. return self.data[self._resolve_source_alias(aliased_item[0]), aliased_item[1]] elif isinstance(aliased_item, str): # Source or key alias. return self.data[self._resolve_any_alias(aliased_item)] raise TypeError('expected alias or (source alias, key) tuple') def _ipython_key_completions_(self): return list(self.data._aliases.keys()) def __contains__(self, aliased_item): try: self[aliased_item] return True except KeyError: return False def __repr__(self): """ Pretty-print all the aliases. """ RED = "\033[91m" END_COLOR = "\033[0m" # Get the right icon for an alias def alias_icon(exists): if isinstance(exists, str): exists = exists in self return " " if exists else f"{RED}✗{END_COLOR}" # Find the alias for a source, if one exists def source_alias(source): for alias, alias_ident in self.data._aliases.items(): if isinstance(alias_ident, str) and source == alias_ident: return alias return None # Group all the aliases by source. The keys of this # dictionary can be either just the source name, or a # tuple of (alias, source). The values are a list of # tuples of (alias, key). source_key_aliases = defaultdict(list) for alias in self.data._aliases.keys(): alias_ident = self.data._aliases[alias] if isinstance(alias_ident, tuple): source = alias_ident[0] if source_alias(source) is not None: dict_key = (source_alias(source), source) else: dict_key = source source_key_aliases[dict_key].append((alias, alias_ident[1])) elif isinstance(alias_ident, str): source_key_aliases[(alias, alias_ident)].extend([]) if len(source_key_aliases) == 0: return "No aliases have been loaded." # Print the aliases output_lines = ["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 alias, source = source output_lines.append(f"{alias_icon(alias)} {alias}: {source}") else: # Check if all the key aliases for the source are valid, # and use that to select an icon for the source keys_exists = [alias in self for alias, _ in alias_keys] if all(keys_exists): source_icon = alias_icon(True) elif not any(keys_exists): source_icon = alias_icon(False) else: source_icon = "~" # Extract the source alias, if it exists if isinstance(source, tuple): source_alias = source[0] source = source[1] else: source_alias = None # If a source has a single key alias, print it on one # line. Otherwise we print the keys indented under the source. if len(alias_keys) == 1: alias, key = alias_keys[0] output_lines.append(f"{alias_icon(alias)} {alias}: ({source}, {key})") else: # If there's an alias, include it in the source header if source_alias is None: source_str = f"{source}" else: source_str = f"{source_alias} ({source})" output_lines.append(f"{source_icon} {source_str}:") for alias, key in alias_keys: output_lines.append(f" {alias_icon(alias)} {alias}: {key}") # Add a newline to the last line added. We can't add a newline by # itself because otherwise it would double up with other newlines # when being joined together at the end output_lines[-1] = output_lines[-1] + "\n" return "\n".join(output_lines) def __str__(self): return f"" def _resolve_aliased_selection(self, selection): if isinstance(selection, dict): res = {self._resolve_source_alias(alias): keys for alias, keys in selection.items()} elif isinstance(selection, Iterable): res = [] for item in selection: if isinstance(item, tuple) and len(item) == 2: # Source alias and literal key. item = (self._resolve_source_alias(item[0]), item[1]) elif isinstance(item, str): item = self._resolve_any_alias(item) if isinstance(item, str): # Source alias. item = (item, '*') res.append(item) return res 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. This method is only accessible through the :attr:`DataCollection.alias` property. In contrast to :meth:`DataCollection.select`, only a subset of ways to select data via aliases is supported: 1. With a source alias and literal key glob pattern:: # Select all pulse energy keys for an aliased XGM fast data. sel = run.alias.select('sa1-xgm', 'data.intensity*') 2. With an iterable of aliases and/or (source alias, key pattern) tuples:: # Select specific keys for an aliased XGM fast data. sel = run.alias.select([('sa1-xgm', 'data.intensitySa1TD'), ('sa1-xgm', 'data.intensitySa3TD')] # Select several aliases, may be both source and key aliases. sel = run.alias.select(['sa1-xgm', 'mono-hv']) Data is included if it matches any of the aliases. Note that this method does not support glob patterns for the source alias. 3. With a dict of source aliases mapped to sets of key names (or empty sets to get all keys):: # Select image.data from an aliased AGIPD and all data # from an aliased XGM. sel = run.select({'agipd': {'image.data'}, 'sa1-xgm': set()}) The optional `require_all` and `require_any` arguments restrict the trains to those for which all or at least one selected sources and keys have at least one data entry. By default, all trains remain selected. Returns a new :class:`DataCollection` object for the selected data. """ if isinstance(seln_or_alias, str): seln_or_alias = [(seln_or_alias, key_glob)] return self.data.select(self._resolve_aliased_selection( seln_or_alias), require_all=require_all, require_any=require_any) def deselect(self, seln_or_alias, key_glob='*'): """Select everything except the specified sources and keys using aliases. This method is only accessible through the :attr:`DataCollection.alias` property. This takes the same arguments as :meth:`select`, but the sources and keys you specify are dropped from the selection. Returns a new :class:`DataCollection` object for the remaining data. """ if isinstance(seln_or_alias, str): seln_or_alias = [(seln_or_alias, key_glob)] return self.data.deselect(self._resolve_aliased_selection( seln_or_alias))