+ 
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+
+"""Logging tools for ess.reduce."""
+
+import logging
+
+
+
+[docs]
+def get_logger() -> logging.Logger:
+ """Return the logger for ess.reduce.
+
+ Returns
+ -------
+ :
+ The requested logger.
+ """
+ return logging.getLogger('scipp.ess.reduce')
+
+
+ 
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+"""Generators for "JSON" data from a NeXus file, for the purpose of testing."""
+
+from collections.abc import Generator
+from typing import Any
+
+import scipp as sc
+
+
+def _as_str_attr(value: Any, name: str) -> dict:
+ val = str(value)
+ return {"string_size": len(val), "type": "string", "name": name, "values": val}
+
+
+def _variable_to_json(var: sc.Variable, name: str):
+ attrs = []
+ if var.dtype == sc.DType.datetime64:
+ offset = var.min()
+ attrs.append(_as_str_attr(offset.value, "offset"))
+ var = var - offset
+ if var.unit is not None:
+ attrs.append(_as_str_attr(var.unit, "units"))
+ return {
+ "module": "dataset",
+ "config": {
+ "name": name,
+ "values": var.values,
+ "size": list(var.shape),
+ "type": str(var.dtype),
+ },
+ "attributes": attrs,
+ }
+
+
+_event_index_0 = sc.array(dims=('dummy',), values=[0], unit=None)
+
+
+def _event_data_pulse_to_json(pulse: sc.DataArray) -> dict:
+ content = pulse.value.coords
+ event_time_zero = sc.concat([pulse.coords['event_time_zero']], 'dummy')
+ event_time_offset = content['event_time_offset']
+ event_id = content.get('event_id')
+ # I think we always have a pixel_id in the flatbuffer, so monitors just get ones?
+ if event_id is None:
+ event_id = sc.ones(sizes=event_time_offset.sizes, dtype='int32', unit=None)
+ children = [
+ _variable_to_json(event_time_zero, name='event_time_zero'),
+ _variable_to_json(event_time_offset, name='event_time_offset'),
+ _variable_to_json(event_id, name='event_id'),
+ _variable_to_json(_event_index_0, name='event_index'),
+ ]
+ group = {
+ "type": "group",
+ "name": "events_0",
+ "children": children,
+ "attributes": [_as_str_attr("NXevent_data", name="NX_class")],
+ }
+ return group
+
+
+
+[docs]
+def event_data_generator(data: sc.DataArray) -> Generator[dict, None, None]:
+ """
+ Generate JSON data for event data from a NeXus file.
+
+ Parameters
+ ----------
+ data:
+ A data array with event data, equivalent to what ScippNexus would load from an
+ NXevent_data group in a NeXus file.
+
+ Yields
+ ------
+ :
+ A dict of data for a single event data pulse that can be wrapped in a
+ :py:class:`ess.reduce.nexus.json_nexus.JSONGroup`.
+ """
+ for pulse in data:
+ yield _event_data_pulse_to_json(pulse)
+
+
+
+
+
+
+"""Adapter for loading NeXus files encoded as JSON."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import numpy as np
+import scippnexus as snx
+
+_nexus_class = "NX_class"
+_nexus_units = "units"
+_nexus_name = "name"
+_nexus_path = "path"
+_nexus_values = "values"
+_nexus_dataset = "dataset"
+_nexus_config = "config"
+_nexus_group = "group"
+_nexus_children = "children"
+_nexus_link = "link"
+_nexus_stream = "stream"
+
+_filewriter_to_supported_numpy_dtype = {
+ "float32": np.float32,
+ "float64": np.float64,
+ "float": np.float32,
+ "double": np.float64,
+ "int8": np.int32,
+ "int16": np.int32,
+ "int32": np.int32,
+ "int64": np.int64,
+ "uint8": np.int32,
+ "uint16": np.int32,
+ "uint32": np.int32,
+ "uint64": np.int64,
+ "string": np.str_,
+}
+
+numpy_to_filewriter_type = {
+ np.float32: "float32",
+ np.float64: "float64",
+ np.int8: "int8",
+ np.int16: "int16",
+ np.int32: "int32",
+ np.int64: "int64",
+ np.uint8: "uint8",
+ np.uint16: "uint16",
+ np.uint32: "uint32",
+ np.uint64: "uint64",
+ np.str_: "string",
+ np.object_: "string",
+}
+
+
+
+[docs]
+def json_nexus_group(
+ json_dict: dict[str, Any], *, definitions: dict[str, type] | None = None
+) -> snx.Group:
+ """Parse a JSON dictionary into a NeXus group.
+
+ Parameters
+ ----------
+ json_dict:
+ ``dict`` containing a NeXus structure as JSON.
+ definitions:
+ ScippNexus application definitions.
+ When not given, the default definitions are used.
+
+ Returns
+ -------
+ :
+ A NeXus group that can be used for loading data as if it were
+ loaded from a file with :class:`scippnexus.File`.
+ """
+ return snx.Group(
+ JSONGroup(json_dict),
+ definitions=definitions if definitions is not None else snx.base_definitions(),
+ )
+
+
+
+
+
+
+
+
+[docs]
+def make_json_attr(name: str, value) -> dict:
+ if isinstance(value, str | bytes):
+ attr_info = {"string_size": len(value), "type": "string"}
+ elif isinstance(value, float):
+ attr_info = {"size": 1, "type": "float64"}
+ elif isinstance(value, int):
+ attr_info = {"size": 1, "type": "int64"}
+ elif isinstance(value, list):
+ attr_info = {"size": len(value), "type": "string"}
+ else:
+ attr_info = {
+ "size": value.shape,
+ "type": numpy_to_filewriter_type[value.dtype.type],
+ }
+ name_and_value = {"name": name, "values": value}
+ return {**attr_info, **name_and_value}
+
+
+
+
+[docs]
+def make_json_dataset(name: str, data) -> dict:
+ if isinstance(data, str | bytes):
+ dataset_info = {"string_size": len(data), "type": "string"}
+ elif isinstance(data, float):
+ dataset_info = {"size": 1, "type": "float64"}
+ elif isinstance(data, int):
+ dataset_info = {"size": 1, "type": "int32"}
+ else:
+ dataset_info = {
+ "size": data.shape,
+ "type": numpy_to_filewriter_type[data.dtype.type],
+ }
+ return {
+ 'module': _nexus_dataset,
+ _nexus_config: {
+ **dataset_info,
+ _nexus_name: name,
+ _nexus_values: data,
+ },
+ "attributes": [],
+ }
+
+
+
+def _get_attribute_value(
+ element: dict, attribute_name: str
+) -> str | float | int | list:
+ """
+ attributes can be a dictionary of key-value pairs, or an array
+ of dictionaries with key, value, type, etc
+ """
+ try:
+ attributes = element["attributes"]
+ try:
+ return attributes[attribute_name]
+ except TypeError:
+ for attribute in attributes:
+ if attribute[_nexus_name] == attribute_name:
+ return attribute[_nexus_values]
+ except KeyError:
+ pass
+ raise MissingAttribute
+
+
+def _visitnodes(root: dict):
+ for child in root.get(_nexus_children, ()):
+ yield child
+ yield from _visitnodes(child)
+
+
+def _name(node: dict):
+ if _nexus_name in node:
+ return node[_nexus_name]
+ if _nexus_config in node:
+ return node[_nexus_config][_nexus_name]
+ return ''
+
+
+def _is_group(node: dict):
+ return _nexus_children in node
+
+
+def _is_dataset(node: dict):
+ return node.get('module') == _nexus_dataset
+
+
+def _is_link(node: dict):
+ return node.get('module') == _nexus_link
+
+
+def _is_stream(node: dict):
+ return 'module' in node and not (_is_dataset(node) or _is_link(node))
+
+
+
+[docs]
+def contains_stream(group: JSONGroup) -> bool:
+ """Return True if the group contains a stream object"""
+ return (
+ isinstance(group, JSONGroup)
+ and _nexus_children in group._node
+ and any(map(_is_stream, group._node[_nexus_children]))
+ )
+
+
+
+
+
+
+
+
+
+
+
+
+[docs]
+class JSONAttributeManager(Mapping[str, Any]):
+
+
+
+ def __contains__(self, name):
+ try:
+ self[name]
+ except MissingAttribute:
+ return False
+ return True
+
+ def __getitem__(self, name):
+ return _get_attribute_value(self._node, name)
+
+ def __setitem__(self, name, value):
+ if name in self:
+ raise NotImplementedError("Replacing existing item not implemented yet.")
+ attr = make_json_attr(name, value)
+ self._node['attributes'].append(attr)
+ return self[name]
+
+ def __iter__(self):
+ if (attrs := self._node.get('attributes')) is not None:
+ if isinstance(attrs, dict):
+ yield from attrs
+ else:
+ for item in attrs:
+ yield item[_nexus_name]
+
+ def __len__(self):
+ return sum(1 for _ in self)
+
+
+
+
+ def get_id(self, name) -> JSONAttrID:
+ # TODO This is a hack that only works since this is used only for a single
+ # purpose by scippnexus.NXobject
+ return JSONAttrID()
+
+
+
+
+[docs]
+class JSONNode:
+
+[docs]
+ def __init__(self, node: dict, *, parent=None):
+ self._file = parent.file if parent is not None else self
+ self._parent = self if parent is None else parent
+ self._node = node
+ name = _name(self._node)
+ if parent is None or parent.name == '/':
+ self._name = f'/{name}'
+ else:
+ self._name = f'{parent.name}/{name}'
+
+
+ @property
+ def attrs(self) -> JSONAttributeManager:
+ return JSONAttributeManager(self._node)
+
+ @property
+ def name(self) -> str:
+ return self._name
+
+ @property
+ def file(self):
+ return self._file
+
+ @property
+ def parent(self):
+ return self._parent
+
+
+
+
+[docs]
+class JSONDataset(JSONNode):
+ @property
+ def dtype(self) -> str:
+ try:
+ dtype = self._node[_nexus_config]["type"]
+ except KeyError:
+ if "dtype" not in self._node[_nexus_config] and isinstance(
+ self._node[_nexus_config][_nexus_values], str
+ ):
+ dtype = 'string'
+ else:
+ dtype = self._node[_nexus_config]["dtype"]
+ if dtype == 'string':
+ return np.dtype(str)
+ return np.dtype(dtype)
+
+ @property
+ def ndim(self) -> int:
+ return len(self.shape)
+
+ @property
+ def shape(self):
+ return np.asarray(self._node[_nexus_config][_nexus_values]).shape
+
+ def __getitem__(self, index):
+ return np.asarray(self._node[_nexus_config][_nexus_values])[index]
+
+ def read_direct(self, buf, source_sel):
+ buf[...] = self[source_sel]
+
+ def asstr(self, **ignored):
+ return self
+
+
+
+
+[docs]
+class JSONGroup(JSONNode):
+ def __contains__(self, name: str) -> bool:
+ try:
+ self[name]
+ return True
+ except KeyError:
+ return False
+
+ def keys(self) -> list[str]:
+ if contains_stream(self):
+ return []
+ children = self._node[_nexus_children]
+ return [_name(child) for child in children if not contains_stream(child)]
+
+ def items(self) -> list[tuple[str, JSONNode]]:
+ return [(key, self[key]) for key in self.keys()]
+
+ def _as_group_or_dataset(self, item, parent):
+ if _is_group(item):
+ return JSONGroup(item, parent=parent)
+ return JSONDataset(item, parent=parent)
+
+ def __getitem__(self, name: str) -> JSONDataset | JSONGroup:
+ if name.startswith('/') and name.count('/') == 1:
+ parent = self.file
+ elif '/' in name:
+ parent = self['/'.join(name.split('/')[:-1])]
+ else:
+ parent = self
+
+ for child in parent._node[_nexus_children]:
+ if _name(child) != name.split('/')[-1]:
+ continue
+ if _is_link(child):
+ return self[child[_nexus_config]["target"]]
+ if _is_group(child) or _is_dataset(child):
+ return self._as_group_or_dataset(child, parent)
+
+ raise KeyError(f"Unable to open object (object '{name}' doesn't exist)")
+
+ def __iter__(self):
+ yield from self.keys()
+
+ def visititems(self, callable):
+ def skip(node):
+ return _is_link(node) or contains_stream(self)
+
+ children = [
+ _name(child) for child in self._node[_nexus_children] if not skip(child)
+ ]
+ for key in children:
+ item = self[key]
+ callable(key, item)
+ if isinstance(item, JSONGroup):
+ item.visititems(callable)
+
+ def create_dataset(self, name: str, data) -> JSONDataset:
+ if name in self:
+ raise NotImplementedError("Replacing existing item not implemented yet.")
+ dataset = make_json_dataset(name, data)
+ self._node[_nexus_children].append(dataset)
+ return self[name]
+
+ def create_group(self, name: str) -> JSONGroup:
+ if name in self:
+ raise NotImplementedError("Replacing existing item not implemented yet.")
+ group = {"type": "group", "name": name, "children": [], "attributes": []}
+ self._node[_nexus_children].append(group)
+ return self[name]
+
+
+
+
+
+
+"""Domain types for use with Sciline, parametrized by run- and monitor-type."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, BinaryIO, Generic, NewType, TypeVar
+
+import sciline
+import scipp as sc
+import scippnexus as snx
+
+FilePath = NewType('FilePath', Path)
+"""Full path to a NeXus file on disk."""
+NeXusFile = NewType('NeXusFile', BinaryIO)
+"""An open NeXus file.
+
+Can be any file handle for reading binary data.
+
+Note that this cannot be used as a parameter in Sciline as there are no
+concrete implementations of ``BinaryIO``.
+The type alias is provided for callers of load functions outside of pipelines.
+"""
+NeXusGroup = NewType('NeXusGroup', snx.Group)
+"""A ScippNexus group in an open file."""
+
+NeXusEntryName = NewType('NeXusEntryName', str)
+"""Name of an entry in a NeXus file."""
+NeXusSourceName = NewType('NeXusSourceName', str)
+"""Name of a source in a NeXus file."""
+
+DetectorBankSizes = NewType("DetectorBankSizes", dict[str, dict[str, int | Any]])
+
+GravityVector = NewType('GravityVector', sc.Variable)
+
+PreopenNeXusFile = NewType('PreopenNeXusFile', bool)
+"""Whether to preopen NeXus files before passing them to the rest of the workflow."""
+
+
+# 1 TypeVars used to parametrize the generic parts of the workflow
+
+# 1.1 Run types
+BackgroundRun = NewType('BackgroundRun', int)
+"""Background run such as a run with only a solvent which the sample is placed in."""
+EmptyBeamRun = NewType('EmptyBeamRun', int)
+"""
+Run with empty sample holder, sometimes called 'direct run'.
+
+It is used for reading the data from the transmission monitor.
+"""
+SampleRun = NewType('SampleRun', int)
+"""Sample run."""
+VanadiumRun = NewType('VanadiumRun', int)
+"""Vanadium run."""
+
+ScatteringRunType = TypeVar(
+ 'ScatteringRunType',
+ BackgroundRun,
+ SampleRun,
+ VanadiumRun,
+)
+
+
+
+[docs]
+class TransmissionRun(Generic[ScatteringRunType]):
+ """
+ Mapping between ScatteringRunType and transmission run.
+
+ In the case where no transmission run is provided, the transmission run should be
+ the same as the measurement (sample or background) run.
+ """
+
+
+
+RunType = TypeVar(
+ 'RunType',
+ BackgroundRun,
+ EmptyBeamRun,
+ SampleRun,
+ # Note that mypy does not seem to like this nesting, may need to find a workaround
+ TransmissionRun[SampleRun],
+ TransmissionRun[BackgroundRun],
+ VanadiumRun,
+)
+"""TypeVar used for specifying BackgroundRun, EmptyBeamRun or SampleRun"""
+
+# 1.2 Monitor types
+Monitor1 = NewType('Monitor1', int)
+"""Identifier for an arbitrary monitor"""
+Monitor2 = NewType('Monitor2', int)
+"""Identifier for an arbitrary monitor"""
+Monitor3 = NewType('Monitor3', int)
+"""Identifier for an arbitrary monitor"""
+Monitor4 = NewType('Monitor4', int)
+"""Identifier for an arbitrary monitor"""
+Monitor5 = NewType('Monitor5', int)
+"""Identifier for an arbitrary monitor"""
+IncidentMonitor = NewType('IncidentMonitor', int)
+"""Incident monitor"""
+TransmissionMonitor = NewType('TransmissionMonitor', int)
+"""Transmission monitor"""
+MonitorType = TypeVar(
+ 'MonitorType',
+ Monitor1,
+ Monitor2,
+ Monitor3,
+ Monitor4,
+ Monitor5,
+ IncidentMonitor,
+ TransmissionMonitor,
+)
+"""TypeVar used for specifying the monitor type such as Incident or Transmission"""
+
+Component = TypeVar(
+ 'Component',
+ snx.NXdetector,
+ snx.NXsample,
+ snx.NXsource,
+ Monitor1,
+ Monitor2,
+ Monitor3,
+ Monitor4,
+ Monitor5,
+ IncidentMonitor,
+ TransmissionMonitor,
+)
+UniqueComponent = TypeVar('UniqueComponent', snx.NXsample, snx.NXsource)
+"""Components that can be identified by their type as there will only be one."""
+
+
+
+[docs]
+class NeXusName(sciline.Scope[Component, str], str):
+ """Name of a component in a NeXus file."""
+
+
+
+
+[docs]
+class NeXusClass(sciline.Scope[Component, str], str):
+ """NX_class of a component in a NeXus file."""
+
+
+
+NeXusDetectorName = NeXusName[snx.NXdetector]
+"""Name of a detector (bank) in a NeXus file."""
+
+
+
+[docs]
+class NeXusComponent(
+ sciline.ScopeTwoParams[Component, RunType, sc.DataGroup], sc.DataGroup
+):
+ """Raw data from a NeXus component."""
+
+
+
+
+[docs]
+class NeXusData(sciline.ScopeTwoParams[Component, RunType, sc.DataArray], sc.DataArray):
+ """
+ Data array loaded from an NXevent_data or NXdata group.
+
+ This must be contained in an NXmonitor or NXdetector group.
+ """
+
+
+
+
+[docs]
+class Position(sciline.ScopeTwoParams[Component, RunType, sc.Variable], sc.Variable):
+ """Position of a component such as source, sample, monitor, or detector."""
+
+
+
+
+[docs]
+class DetectorPositionOffset(sciline.Scope[RunType, sc.Variable], sc.Variable):
+ """Offset for the detector position, added to base position."""
+
+
+
+
+[docs]
+class MonitorPositionOffset(
+ sciline.ScopeTwoParams[RunType, MonitorType, sc.Variable], sc.Variable
+):
+ """Offset for the monitor position, added to base position."""
+
+
+
+
+[docs]
+class CalibratedDetector(sciline.Scope[RunType, sc.DataArray], sc.DataArray):
+ """Calibrated data from a detector."""
+
+
+
+
+[docs]
+class CalibratedBeamline(sciline.Scope[RunType, sc.DataArray], sc.DataArray):
+ """Calibrated beamline with detector and other components."""
+
+
+
+
+[docs]
+class CalibratedMonitor(
+ sciline.ScopeTwoParams[RunType, MonitorType, sc.DataArray], sc.DataArray
+):
+ """Calibrated data from a monitor."""
+
+
+
+
+[docs]
+class DetectorData(sciline.Scope[RunType, sc.DataArray], sc.DataArray):
+ """Calibrated detector merged with neutron event or histogram data."""
+
+
+
+
+[docs]
+class MonitorData(
+ sciline.ScopeTwoParams[RunType, MonitorType, sc.DataArray], sc.DataArray
+):
+ """Calibrated monitor merged with neutron event or histogram data."""
+
+
+
+
+
+
+
+
+[docs]
+@dataclass
+class TimeInterval(Generic[RunType]):
+ """Range of neutron pulses to load from NXevent_data or NXdata groups."""
+
+ value: slice
+
+
+
+
+[docs]
+@dataclass
+class NeXusFileSpec(Generic[RunType]):
+ value: FilePath | NeXusFile | NeXusGroup
+
+
+
+
+[docs]
+@dataclass
+class NeXusLocationSpec:
+ """
+ NeXus filename and optional parameters to identify (parts of) a component to load.
+ """
+
+ filename: FilePath | NeXusFile | NeXusGroup
+ entry_name: NeXusEntryName | None = None
+ component_name: str | None = None
+ selection: snx.typing.ScippIndex | slice = ()
+
+
+
+
+[docs]
+@dataclass
+class NeXusComponentLocationSpec(NeXusLocationSpec, Generic[Component, RunType]):
+ """
+ NeXus filename and optional parameters to identify (parts of) a component to load.
+ """
+
+
+
+
+[docs]
+@dataclass
+class NeXusDataLocationSpec(NeXusLocationSpec, Generic[Component, RunType]):
+ """NeXus filename and parameters to identify (parts of) detector data to load."""
+
+
+
+T = TypeVar('T', bound='NeXusTransformationChain')
+
+
+
+
+ sciline.ScopeTwoParams[Component, RunType, snx.TransformationChain],
+ snx.TransformationChain,
+): ...
+
+
+
+[docs]
+@dataclass
+class NeXusTransformation(Generic[Component, RunType]):
+ value: sc.Variable
+
+
+[docs]
+ @staticmethod
+ def from_chain(
+ chain: NeXusTransformationChain[Component, RunType],
+ ) -> 'NeXusTransformation[Component, RunType]':
+ """
+ Convert a transformation chain to a single transformation.
+
+ As transformation chains may be time-dependent, this method will need to select
+ a specific time point to convert to a single transformation. This may include
+ averaging as well as threshold checks. This is not implemented yet and we
+ therefore currently raise an error if the transformation chain does not compute
+ to a scalar.
+ """
+ if chain.transformations.sizes != {}:
+ raise ValueError(f"Expected scalar transformation, got {chain}")
+ transform = chain.compute()
+ return NeXusTransformation(value=transform)
+
+
+
+
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+
+"""Workflow and workflow components for interacting with NeXus files."""
+
+from collections.abc import Sequence
+from copy import deepcopy
+from typing import Any
+
+import networkx as nx
+import sciline
+import scipp as sc
+import scippnexus as snx
+from scipp.constants import g
+from scipp.core import label_based_index_to_positional_index
+
+from . import _nexus_loader as nexus
+from .types import (
+ CalibratedBeamline,
+ CalibratedDetector,
+ CalibratedMonitor,
+ Component,
+ DetectorBankSizes,
+ DetectorData,
+ DetectorPositionOffset,
+ Filename,
+ GravityVector,
+ MonitorData,
+ MonitorPositionOffset,
+ MonitorType,
+ NeXusClass,
+ NeXusComponent,
+ NeXusComponentLocationSpec,
+ NeXusData,
+ NeXusDataLocationSpec,
+ NeXusFileSpec,
+ NeXusName,
+ NeXusTransformation,
+ NeXusTransformationChain,
+ Position,
+ PreopenNeXusFile,
+ RunType,
+ TimeInterval,
+ UniqueComponent,
+)
+
+origin = sc.vector([0, 0, 0], unit="m")
+"""The origin, used as default sample position."""
+no_offset = sc.vector([0, 0, 0], unit="m")
+"""Offset that does not change the position."""
+
+
+
+[docs]
+def file_path_to_file_spec(
+ filename: Filename[RunType], preopen: PreopenNeXusFile
+) -> NeXusFileSpec[RunType]:
+ return NeXusFileSpec[RunType](
+ snx.File(filename, definitions=definitions) if preopen else filename
+ )
+
+
+
+
+[docs]
+def no_monitor_position_offset() -> MonitorPositionOffset[RunType, MonitorType]:
+ return MonitorPositionOffset[RunType, MonitorType](no_offset)
+
+
+
+
+[docs]
+def no_detector_position_offset() -> DetectorPositionOffset[RunType]:
+ return DetectorPositionOffset[RunType](no_offset)
+
+
+
+
+[docs]
+def full_time_interval() -> TimeInterval[RunType]:
+ """Select all neutron pulses in the data."""
+ return TimeInterval[RunType](slice(None, None))
+
+
+
+
+[docs]
+def gravity_vector_neg_y() -> GravityVector:
+ """
+ Gravity vector for default instrument coordinate system where y is up.
+ """
+ return GravityVector(sc.vector(value=[0, -1, 0]) * g)
+
+
+
+
+[docs]
+def component_spec_by_name(
+ filename: NeXusFileSpec[RunType], name: NeXusName[Component]
+) -> NeXusComponentLocationSpec[Component, RunType]:
+ """
+ Create a location spec for a component group in a NeXus file.
+
+ Parameters
+ ----------
+ filename:
+ NeXus file to use for the location spec.
+ name:
+ Name of the component group.
+ """
+ return NeXusComponentLocationSpec[Component, RunType](
+ filename=filename.value, component_name=name
+ )
+
+
+
+
+[docs]
+def unique_component_spec(
+ filename: NeXusFileSpec[RunType],
+) -> NeXusComponentLocationSpec[UniqueComponent, RunType]:
+ """
+ Create a location spec for a unique component group in a NeXus file.
+
+ Parameters
+ ----------
+ filename:
+ NeXus file to use for the location spec.
+ """
+ return NeXusComponentLocationSpec[UniqueComponent, RunType](filename=filename.value)
+
+
+
+
+[docs]
+def data_by_name(
+ filename: NeXusFileSpec[RunType],
+ name: NeXusName[Component],
+ selection: TimeInterval[RunType],
+) -> NeXusDataLocationSpec[Component, RunType]:
+ """
+ Create a location spec for monitor or detector data in a NeXus file.
+
+ Parameters
+ ----------
+ filename:
+ NeXus file to use for the location spec.
+ name:
+ Name of the monitor or detector group.
+ selection:
+ Time range (start and stop as a Python slice object).
+ """
+ return NeXusDataLocationSpec[Component, RunType](
+ filename=filename.value, component_name=name, selection=selection.value
+ )
+
+
+
+
+[docs]
+def load_nexus_sample(
+ location: NeXusComponentLocationSpec[snx.NXsample, RunType],
+) -> NeXusComponent[snx.NXsample, RunType]:
+ """
+ Load a NeXus sample group from a file.
+
+ If there is no sample group in the file, an empty group is returned. This should
+ not happen, but handling it gracefully makes testing and working with
+ pre-production files easier. There should be little harm in returning an empty
+ group. Subsequent extract of the sample position will then default to the origin.
+
+ Parameters
+ ----------
+ location:
+ Location spec for the sample group.
+ """
+ try:
+ dg = nexus.load_component(location, nx_class=snx.NXsample)
+ except ValueError:
+ dg = sc.DataGroup()
+ if 'depends_on' not in dg:
+ dg['depends_on'] = snx.TransformationChain(parent='', value='.')
+ return NeXusComponent[snx.NXsample, RunType](dg)
+
+
+
+
+[docs]
+def nx_class_for_monitor() -> NeXusClass[MonitorType]:
+ return NeXusClass[MonitorType](snx.NXmonitor)
+
+
+
+
+[docs]
+def nx_class_for_detector() -> NeXusClass[snx.NXdetector]:
+ return NeXusClass[snx.NXdetector](snx.NXdetector)
+
+
+
+
+[docs]
+def nx_class_for_source() -> NeXusClass[snx.NXsource]:
+ return NeXusClass[snx.NXsource](snx.NXsource)
+
+
+
+
+[docs]
+def nx_class_for_sample() -> NeXusClass[snx.NXsample]:
+ return NeXusClass[snx.NXsample](snx.NXsample)
+
+
+
+
+[docs]
+def load_nexus_component(
+ location: NeXusComponentLocationSpec[Component, RunType],
+ nx_class: NeXusClass[Component],
+) -> NeXusComponent[Component, RunType]:
+ """
+ Load a NeXus component group from a file.
+
+ When loading a detector or monitor, event data is replaced by placeholders.
+
+ As the event data can be large and is not needed at this stage, it is replaced by
+ a placeholder. A placeholder is used to allow for returning a scipp.DataArray, which
+ is what most downstream code will expect.
+ Currently the placeholder is the detector number (for detectors) or a size-0 array
+ (for monitors), but this may change in the future.
+
+ The returned object is a scipp.DataGroup, as it may contain additional information
+ about the detector that cannot be represented as a single scipp.DataArray. Most
+ downstream code will only be interested in the contained scipp.DataArray so this
+ needs to be extracted. However, other processing steps may require the additional
+ information, so it is kept in the DataGroup.
+
+ Parameters
+ ----------
+ location:
+ Location spec for the source group.
+ nx_class:
+ NX_class to identify the component.
+ """
+ return NeXusComponent[Component, RunType](
+ nexus.load_component(location, nx_class=nx_class, definitions=definitions)
+ )
+
+
+
+
+[docs]
+def load_nexus_data(
+ location: NeXusDataLocationSpec[Component, RunType],
+) -> NeXusData[Component, RunType]:
+ """
+ Load event or histogram data from a NeXus detector group.
+
+ Parameters
+ ----------
+ location:
+ Location spec for the detector group.
+ """
+ return NeXusData[Component, RunType](
+ nexus.load_data(
+ file_path=location.filename,
+ entry_name=location.entry_name,
+ selection=location.selection,
+ component_name=location.component_name,
+ )
+ )
+
+
+
+
+[docs]
+def get_transformation_chain(
+ detector: NeXusComponent[Component, RunType],
+) -> NeXusTransformationChain[Component, RunType]:
+ """
+ Extract the transformation chain from a NeXus detector group.
+
+ Parameters
+ ----------
+ detector:
+ NeXus detector group.
+ """
+ chain = detector['depends_on']
+ return NeXusTransformationChain[Component, RunType](chain)
+
+
+
+def _time_filter(transform: sc.DataArray) -> sc.Variable:
+ if transform.ndim == 0 or transform.sizes == {'time': 1}:
+ return transform.data.squeeze()
+ raise ValueError(
+ f"Transform is time-dependent: {transform}, but no filter is provided."
+ )
+
+
+
+[docs]
+def to_transformation(
+ chain: NeXusTransformationChain[Component, RunType], interval: TimeInterval[RunType]
+) -> NeXusTransformation[Component, RunType]:
+ """
+ Convert transformation chain into a single transformation matrix.
+
+ If one or more transformations in the chain are time-dependent, the time interval
+ is used to select a specific time point. If the interval is not a single time point,
+ an error is raised. This may be extended in the future to a more sophisticated
+ mechanism, e.g., averaging over the interval to remove noise.
+
+ Parameters
+ ----------
+ chain:
+ Transformation chain.
+ interval:
+ Time interval to select from the transformation chain.
+ """
+
+ chain = deepcopy(chain)
+ for t in chain.transformations.values():
+ if t.sizes == {} or not isinstance(t.value, sc.DataArray):
+ continue
+ start = interval.value.start
+ stop = interval.value.stop
+ if isinstance(start, sc.Variable) or isinstance(stop, sc.Variable):
+ # NXlog entries are generally interpreted as the previous value being valid
+ # until the next entry. We therefore need to select the previous value, and
+ # any index after the last entry refers to the last entry, i.e., there is no
+ # "end" time in the files. We add a dummy end so we can use Scipp's label-
+ # based indexing for histogram data.
+ time = t.value.coords['time']
+ # Add 1000 days as a dummy end time. This is hopefully long enough to cover
+ # all reasonable use cases.
+ delta = sc.scalar(24_000, unit='hours', dtype='int64').to(unit=time.unit)
+ time = sc.concat([time, time[-1] + delta], 'time')
+ idx = label_based_index_to_positional_index(
+ sizes=t.sizes, coord=time, index=interval.value
+ )
+ t.value = _time_filter(t.value[idx])
+ else:
+ t.value = _time_filter(t.value['time', interval.value])
+
+ return NeXusTransformation[Component, RunType].from_chain(chain)
+
+
+
+
+[docs]
+def compute_position(
+ transformation: NeXusTransformation[Component, RunType],
+) -> Position[Component, RunType]:
+ """Compute the position of a component from a transformation matrix."""
+ return Position[Component, RunType](transformation.value * origin)
+
+
+
+
+[docs]
+def get_calibrated_detector(
+ detector: NeXusComponent[snx.NXdetector, RunType],
+ *,
+ transform: NeXusTransformation[snx.NXdetector, RunType],
+ # Strictly speaking we could apply an offset by modifying the transformation chain,
+ # using a more generic implementation. However, this may in general require
+ # extending the chain and it is currently not clear if that is desirable. As far as
+ # I am aware the offset is currently mainly used for handling files from other
+ # facilities and it is not clear if it is needed for ESS data and should be kept at
+ # all.
+ offset: DetectorPositionOffset[RunType],
+ bank_sizes: DetectorBankSizes,
+) -> CalibratedDetector[RunType]:
+ """
+ Extract the data array corresponding to a detector's signal field.
+
+ The returned data array includes coords and masks pertaining directly to the
+ signal values array, but not additional information about the detector. The
+ data array is reshaped to the logical detector shape, which by folding the data
+ array along the detector_number dimension.
+
+ Parameters
+ ----------
+ detector:
+ NeXus detector group.
+ offset:
+ Offset to add to the detector position.
+ bank_sizes:
+ Dictionary of detector bank sizes.
+ """
+ da = nexus.extract_signal_data_array(detector)
+ if (
+ sizes := (bank_sizes or {}).get(detector.get('nexus_component_name'))
+ ) is not None:
+ da = da.fold(dim="detector_number", sizes=sizes)
+ # Note: We apply offset as early as possible, i.e., right in this function
+ # the detector array from the raw loader NeXus group, to prevent a source of bugs.
+ # If the NXdetector in the file is not 1-D, we want to match the order of dims.
+ # zip_pixel_offsets otherwise yields a vector with dimensions in the order given
+ # by the x/y/z offsets.
+ offsets = snx.zip_pixel_offsets(da.coords).transpose(da.dims).copy()
+ # We use the unit of the offsets as this is likely what the user expects.
+ if transform.value.unit is not None and transform.value.unit != '':
+ transform_value = transform.value.to(unit=offsets.unit)
+ else:
+ transform_value = transform.value
+ position = transform_value * offsets
+ return CalibratedDetector[RunType](
+ da.assign_coords(position=position + offset.to(unit=position.unit))
+ )
+
+
+
+
+[docs]
+def assemble_beamline(
+ detector: CalibratedDetector[RunType],
+ source_position: Position[snx.NXsource, RunType],
+ sample_position: Position[snx.NXsample, RunType],
+ gravity: GravityVector,
+) -> CalibratedBeamline[RunType]:
+ """
+ Add beamline information (gravity vector, source- and sample-position) to detector.
+
+ This is performed separately and after :py:func:`get_calibrated_detector` to avoid
+ as false dependency of, e.g., the reshaped detector numbers on the sample position.
+ The latter can change during a run, e.g., for a rotating sample. The detector
+ numbers might be used, e.g., to mask certain detector pixels, and should not depend
+ on the sample position.
+
+ Parameters
+ ----------
+ detector:
+ NeXus detector group.
+ source_position:
+ Position of the neutron source.
+ sample_position:
+ Position of the sample.
+ gravity:
+ Gravity vector.
+ """
+ return CalibratedBeamline[RunType](
+ detector.assign_coords(
+ source_position=source_position,
+ sample_position=sample_position,
+ gravity=gravity,
+ )
+ )
+
+
+
+
+[docs]
+def assemble_detector_data(
+ detector: CalibratedBeamline[RunType],
+ event_data: NeXusData[snx.NXdetector, RunType],
+) -> DetectorData[RunType]:
+ """
+ Assemble a detector data array with event data.
+
+ Also adds variances to the event data if they are missing.
+
+ Parameters
+ ----------
+ detector:
+ Calibrated detector data array.
+ event_data:
+ Event data array.
+ """
+ grouped = nexus.group_event_data(
+ event_data=event_data, detector_number=detector.coords['detector_number']
+ )
+ return DetectorData[RunType](
+ _add_variances(grouped)
+ .assign_coords(detector.coords)
+ .assign_masks(detector.masks)
+ )
+
+
+
+
+[docs]
+def get_calibrated_monitor(
+ monitor: NeXusComponent[MonitorType, RunType],
+ offset: MonitorPositionOffset[RunType, MonitorType],
+ source_position: Position[snx.NXsource, RunType],
+) -> CalibratedMonitor[RunType, MonitorType]:
+ """
+ Extract the data array corresponding to a monitor's signal field.
+
+ The returned data array includes coords pertaining directly to the
+ signal values array, but not additional information about the monitor.
+
+ Parameters
+ ----------
+ monitor:
+ NeXus monitor group.
+ offset:
+ Offset to add to the monitor position.
+ source_position:
+ Position of the neutron source.
+ """
+ monitor = nexus.compute_component_position(monitor)
+ return CalibratedMonitor[RunType, MonitorType](
+ nexus.extract_signal_data_array(monitor).assign_coords(
+ position=monitor['position'] + offset.to(unit=monitor['position'].unit),
+ source_position=source_position,
+ )
+ )
+
+
+
+
+[docs]
+def assemble_monitor_data(
+ monitor: CalibratedMonitor[RunType, MonitorType],
+ data: NeXusData[MonitorType, RunType],
+) -> MonitorData[RunType, MonitorType]:
+ """
+ Assemble a monitor data array with event data.
+
+ Also adds variances to the event data if they are missing.
+
+ Parameters
+ ----------
+ monitor:
+ Calibrated monitor data array.
+ data:
+ Data array with neutron counts.
+ """
+ da = data.assign_coords(monitor.coords).assign_masks(monitor.masks)
+ return MonitorData[RunType, MonitorType](_add_variances(da))
+
+
+
+def _drop(
+ children: dict[str, snx.Field | snx.Group], classes: tuple[snx.NXobject, ...]
+) -> dict[str, snx.Field | snx.Group]:
+ return {
+ name: child
+ for name, child in children.items()
+ if not (isinstance(child, snx.Group) and (child.nx_class in classes))
+ }
+
+
+class _StrippedDetector(snx.NXdetector):
+ """Detector definition without large geometry or event data for ScippNexus.
+
+ Drops NXoff_geometry and NXevent_data groups, data is replaced by detector_number.
+ """
+
+ def __init__(
+ self, attrs: dict[str, Any], children: dict[str, snx.Field | snx.Group]
+ ):
+ children = _drop(children, (snx.NXoff_geometry, snx.NXevent_data))
+ children['data'] = children['detector_number']
+ super().__init__(attrs=attrs, children=children)
+
+
+class _DummyField:
+ """Dummy field that can replace snx.Field in NXmonitor."""
+
+ def __init__(self):
+ self.attrs = {}
+ self.sizes = {'event_time_zero': 0}
+ self.dims = ('event_time_zero',)
+ self.shape = (0,)
+
+ def __getitem__(self, key: Any) -> sc.Variable:
+ return sc.empty(dims=self.dims, shape=self.shape, unit=None)
+
+
+class _StrippedMonitor(snx.NXmonitor):
+ """Monitor definition without event data for ScippNexus.
+
+ Drops NXevent_data group, data is replaced by a dummy field.
+ """
+
+ def __init__(
+ self, attrs: dict[str, Any], children: dict[str, snx.Field | snx.Group]
+ ):
+ children = _drop(children, (snx.NXevent_data,))
+ children['data'] = _DummyField()
+ super().__init__(attrs=attrs, children=children)
+
+
+def _add_variances(da: sc.DataArray) -> sc.DataArray:
+ out = da.copy(deep=False)
+ if out.bins is not None:
+ content = out.bins.constituents['data']
+ if content.variances is None:
+ content.variances = content.values
+ return out
+
+
+definitions = snx.base_definitions()
+definitions["NXdetector"] = _StrippedDetector
+definitions["NXmonitor"] = _StrippedMonitor
+
+
+_common_providers = (
+ gravity_vector_neg_y,
+ file_path_to_file_spec,
+ full_time_interval,
+ component_spec_by_name,
+ unique_component_spec, # after component_spec_by_name, partially overrides
+ get_transformation_chain,
+ to_transformation,
+ compute_position,
+ load_nexus_data,
+ load_nexus_component,
+ data_by_name,
+ nx_class_for_detector,
+ nx_class_for_monitor,
+ nx_class_for_source,
+ nx_class_for_sample,
+)
+
+_monitor_providers = (
+ no_monitor_position_offset,
+ get_calibrated_monitor,
+ assemble_monitor_data,
+)
+
+_detector_providers = (
+ no_detector_position_offset,
+ load_nexus_sample,
+ get_calibrated_detector,
+ assemble_beamline,
+ assemble_detector_data,
+)
+
+
+
+[docs]
+def LoadMonitorWorkflow() -> sciline.Pipeline:
+ """Generic workflow for loading monitor data from a NeXus file."""
+ wf = sciline.Pipeline((*_common_providers, *_monitor_providers))
+ wf[PreopenNeXusFile] = PreopenNeXusFile(False)
+ return wf
+
+
+
+
+[docs]
+def LoadDetectorWorkflow() -> sciline.Pipeline:
+ """Generic workflow for loading detector data from a NeXus file."""
+ wf = sciline.Pipeline((*_common_providers, *_detector_providers))
+ wf[DetectorBankSizes] = DetectorBankSizes({})
+ wf[PreopenNeXusFile] = PreopenNeXusFile(False)
+ return wf
+
+
+
+
+[docs]
+def GenericNeXusWorkflow(
+ *,
+ run_types: Sequence[sciline.typing.Key] | None = None,
+ monitor_types: Sequence[sciline.typing.Key] | None = None,
+) -> sciline.Pipeline:
+ """
+ Generic workflow for loading detector and monitor data from a NeXus file.
+
+ Parameters
+ ----------
+ run_types:
+ List of run types to include in the workflow. If not provided, all run types
+ are included. It is recommended to specify run types to avoid creating very
+ large workflows.
+ monitor_types:
+ List of monitor types to include in the workflow. If not provided, all monitor
+ types are included. It is recommended to specify monitor types to avoid creating
+ very large workflows.
+
+ Returns
+ -------
+ :
+ The workflow.
+ """
+ if monitor_types is not None and run_types is None:
+ raise ValueError("run_types must be specified if monitor_types is specified")
+ wf = sciline.Pipeline(
+ (*_common_providers, *_monitor_providers, *_detector_providers)
+ )
+ wf[DetectorBankSizes] = DetectorBankSizes({})
+ wf[PreopenNeXusFile] = PreopenNeXusFile(False)
+
+ g = wf.underlying_graph
+ ancestors = set()
+ # DetectorData and MonitorData are the "final" outputs, so finding and removing all
+ # their ancestors is what we need to strip unused run and monitor types.
+ for rt in run_types or ():
+ ancestors |= nx.ancestors(g, DetectorData[rt])
+ ancestors.add(DetectorData[rt])
+ for mt in monitor_types or ():
+ ancestors |= nx.ancestors(g, MonitorData[rt, mt])
+ ancestors.add(MonitorData[rt, mt])
+ if run_types is not None:
+ g.remove_nodes_from(set(g.nodes) - ancestors)
+ return wf
+
+
+
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+"""This module provides tools for running workflows in a streaming fashion."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+
+import networkx as nx
+import sciline
+import scipp as sc
+
+T = TypeVar('T')
+
+
+
+[docs]
+def maybe_hist(value: T) -> T:
+ """
+ Convert value to a histogram if it is not already a histogram.
+
+ This is the default pre-processing used by accumulators.
+
+ Parameters
+ ----------
+ value:
+ Value to be converted to a histogram.
+
+ Returns
+ -------
+ :
+ Histogram.
+ """
+ return value if value.bins is None else value.hist()
+
+
+
+
+[docs]
+class Accumulator(ABC, Generic[T]):
+ """
+ Abstract base class for accumulators.
+
+ Accumulators are used to accumulate values over multiple chunks.
+ """
+
+
+[docs]
+ def __init__(self, preprocess: Callable[[T], T] | None = maybe_hist) -> None:
+ """
+ Parameters
+ ----------
+ preprocess:
+ Preprocessing function to be applied to pushed values prior to accumulation.
+ """
+ self._preprocess = preprocess
+
+
+
+[docs]
+ def push(self, value: T) -> None:
+ """
+ Push a value to the accumulator.
+
+ Parameters
+ ----------
+ value:
+ Value to be pushed to the accumulator.
+ """
+ if self._preprocess is not None:
+ value = self._preprocess(value)
+ self._do_push(value)
+
+
+ @abstractmethod
+ def _do_push(self, value: T) -> None: ...
+
+ @property
+ @abstractmethod
+ def value(self) -> T:
+ """
+ Get the accumulated value.
+
+ Returns
+ -------
+ :
+ Accumulated value.
+ """
+
+
+
+
+[docs]
+class EternalAccumulator(Accumulator[T]):
+ """
+ Simple accumulator that adds pushed values immediately.
+
+ Does not support event data.
+ """
+
+
+[docs]
+ def __init__(self, **kwargs: Any) -> None:
+ super().__init__(**kwargs)
+ self._value: T | None = None
+
+
+ @property
+ def value(self) -> T:
+ return self._value.copy()
+
+ def _do_push(self, value: T) -> None:
+ if self._value is None:
+ self._value = value.copy()
+ else:
+ self._value += value
+
+
+
+
+[docs]
+class RollingAccumulator(Accumulator[T]):
+ """
+ Accumulator that adds pushed values to a rolling window.
+
+ Does not support event data.
+ """
+
+
+[docs]
+ def __init__(self, window: int = 10, **kwargs: Any) -> None:
+ """
+ Parameters
+ ----------
+ window:
+ Size of the rolling window.
+ """
+ super().__init__(**kwargs)
+ self._window = window
+ self._values: list[T] = []
+
+
+ @property
+ def value(self) -> T:
+ # Naive and potentially slow implementation if values and/or window are large!
+ return sc.reduce(self._values).sum()
+
+ def _do_push(self, value: T) -> None:
+ self._values.append(value)
+ if len(self._values) > self._window:
+ self._values.pop(0)
+
+
+
+
+[docs]
+class StreamProcessor:
+ """
+ Wrap a base workflow for streaming processing of chunks.
+
+ Note that this class can not determine if the workflow is valid for streamed
+ processing based on the input keys. In particular, it is the responsibility of the
+ user to ensure that the workflow is "linear" with respect to the dynamic keys up to
+ the accumulation keys.
+ """
+
+
+[docs]
+ def __init__(
+ self,
+ base_workflow: sciline.Pipeline,
+ *,
+ dynamic_keys: tuple[sciline.typing.Key, ...],
+ target_keys: tuple[sciline.typing.Key, ...],
+ accumulators: dict[sciline.typing.Key, Accumulator, Callable[..., Accumulator]]
+ | tuple[sciline.typing.Key, ...],
+ ) -> None:
+ """
+ Create a stream processor.
+
+ Parameters
+ ----------
+ base_workflow:
+ Workflow to be used for processing chunks.
+ dynamic_keys:
+ Keys that are expected to be updated with each chunk.
+ target_keys:
+ Keys to be computed and returned.
+ accumulators:
+ Keys at which to accumulate values and their accumulators. If a tuple is
+ passed, :py:class:`EternalAccumulator` is used for all keys. Otherwise, a
+ dict mapping keys to accumulator instances can be passed. If a dict value is
+ a callable, base_workflow.bind_and_call(value) is used to make an instance.
+ """
+ workflow = sciline.Pipeline()
+ for key in target_keys:
+ workflow[key] = base_workflow[key]
+ for key in dynamic_keys:
+ workflow[key] = None # hack to prune branches
+
+ # Find and pre-compute static nodes as far down the graph as possible
+ # See also https://github.com/scipp/sciline/issues/148.
+ nodes = _find_descendants(workflow, dynamic_keys)
+ parents = _find_parents(workflow, nodes) - nodes
+ for key, value in base_workflow.compute(parents).items():
+ workflow[key] = value
+
+ self._process_chunk_workflow = workflow.copy()
+ self._finalize_workflow = workflow.copy()
+ self._accumulators = (
+ accumulators
+ if isinstance(accumulators, dict)
+ else {key: EternalAccumulator() for key in accumulators}
+ )
+ # Depending on the target_keys, some accumulators can be unused and should not
+ # be computed when adding a chunk.
+ self._accumulators = {
+ key: value
+ for key, value in self._accumulators.items()
+ if key in self._process_chunk_workflow.underlying_graph
+ }
+ # Create accumulators unless instances were passed. This allows for initializing
+ # accumulators with arguments that depend on the workflow such as bin edges,
+ # which would otherwise be hard to obtain.
+ self._accumulators = {
+ key: value
+ if isinstance(value, Accumulator)
+ else base_workflow.bind_and_call(value)
+ for key, value in self._accumulators.items()
+ }
+ self._target_keys = target_keys
+
+
+ def add_chunk(
+ self, chunks: dict[sciline.typing.Key, Any]
+ ) -> dict[sciline.typing.Key, Any]:
+ for key, value in chunks.items():
+ self._process_chunk_workflow[key] = value
+ # There can be dynamic keys that do not "terminate" in any accumulator. In
+ # that case, we need to make sure they can be and are used when computing
+ # the target keys.
+ self._finalize_workflow[key] = value
+ to_accumulate = self._process_chunk_workflow.compute(self._accumulators)
+ for key, processed in to_accumulate.items():
+ self._accumulators[key].push(processed)
+ self._finalize_workflow[key] = self._accumulators[key].value
+ return self._finalize_workflow.compute(self._target_keys)
+
+
+
+def _find_descendants(
+ workflow: sciline.Pipeline, keys: tuple[sciline.typing.Key, ...]
+) -> set[sciline.typing.Key]:
+ graph = workflow.underlying_graph
+ descendants = set()
+ for key in keys:
+ descendants |= nx.descendants(graph, key)
+ return descendants | set(keys)
+
+
+def _find_parents(
+ workflow: sciline.Pipeline, keys: tuple[sciline.typing.Key, ...]
+) -> set[sciline.typing.Key]:
+ graph = workflow.underlying_graph
+ parents = set()
+ for key in keys:
+ parents |= set(graph.predecessors(key))
+ return parents
+
+
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+"""Tools for handling statistical uncertainties.
+
+This module provides tools for handling statistical uncertainties in the context of
+data reduction. Handling variances during broadcast operations is not handled
+correctly by Scipp because correlations are not tracked.
+See https://doi.org/10.3233/JNR-220049 for context.
+
+This module provides three ways of handling variances during broadcast operations,
+defined by :py:class:`UncertaintyBroadcastMode`.
+The recommended use of this module is via the :py:func:`broadcast_uncertainties`
+helper function.
+"""
+
+from enum import Enum
+from typing import TypeVar, overload
+
+import numpy as np
+import scipp as sc
+from scipp.core.concepts import irreducible_mask
+
+T = TypeVar("T", bound=sc.Variable | sc.DataArray)
+
+
+UncertaintyBroadcastMode = Enum(
+ 'UncertaintyBroadcastMode', ['drop', 'upper_bound', 'fail']
+)
+"""
+Mode for broadcasting uncertainties.
+
+- `drop`: Drop variances if the data is broadcasted.
+- `upper_bound`: Compute an upper bound for the variances.
+- `fail`: Do not broadcast, simply return the input data.
+
+See https://doi.org/10.3233/JNR-220049 for context.
+"""
+
+
+@overload
+def broadcast_with_upper_bound_variances(
+ data: sc.Variable, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.Variable:
+ pass
+
+
+@overload
+def broadcast_with_upper_bound_variances(
+ data: sc.DataArray, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.DataArray:
+ pass
+
+
+
+[docs]
+def broadcast_with_upper_bound_variances(
+ data: sc.Variable | sc.DataArray, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.Variable | sc.DataArray:
+ """
+ Compute an upper bound for the variances of the broadcasted data.
+
+ The variances of the broadcasted data are computed by scaling the variances of the
+ input data by the volume of the new subspace. The volume of the new subspace is
+ computed as the product of the sizes of the new dimensions. In the case of an
+ event-data prototype the events are counted.
+
+ Parameters
+ ----------
+ data:
+ The data to broadcast.
+ prototype:
+ Defines the new sizes (dims and shape). If present, masks are used to exclude
+ masked values from the variance computation.
+
+ Returns
+ -------
+ :
+ The data with the variances scaled by the volume of the new subspace.
+ """
+ if _no_variance_broadcast(data, prototype=prototype):
+ return data
+ for dim in prototype.dims:
+ coord1 = None if isinstance(data, sc.Variable) else data.coords.get(dim)
+ coord2 = (
+ None if isinstance(prototype, sc.Variable) else prototype.coords.get(dim)
+ )
+ if coord1 is None or coord2 is None:
+ if dim in data.dims:
+ if data.sizes[dim] != prototype.sizes[dim]:
+ raise ValueError("Mismatching binning not supported in broadcast.")
+ continue
+ elif sc.identical(coord1, coord2):
+ continue
+ raise ValueError("Mismatching binning not supported in broadcast.")
+ sizes = prototype.sizes
+ mask = sc.scalar(False)
+ if isinstance(prototype, sc.DataArray):
+ if (irred := irreducible_mask(prototype, dim=sizes)) is not None:
+ for dim in data.dims:
+ if dim in irred.dims:
+ irred = irred.all(dim)
+ mask = irred
+ data = data.copy()
+ sizes = {**sizes, **data.sizes}
+ if prototype.bins is None:
+ size = (~mask).sum().to(dtype='int64', copy=False)
+ for dim, dim_size in sizes.items():
+ if dim not in data.dims and dim not in mask.dims:
+ size *= sc.index(dim_size)
+ else:
+ size = prototype.bins.size().sum(set(prototype.dims) - set(data.dims))
+ scale = size.broadcast(sizes=sizes).to(dtype='float64')
+ if not sc.identical(mask, sc.scalar(False)):
+ # The masked values are not counted in the variance, so we set them to infinity.
+ scale.values[mask.broadcast(sizes=sizes).values] = np.inf
+ data = data.broadcast(sizes=sizes).copy()
+ data.variances *= scale.values
+ if prototype.bins is not None:
+ # Note that we are not using event masks in the upper-bound computation. Less
+ # than optimal, but simpler.
+ if isinstance(data, sc.Variable):
+ data = sc.bins_like(prototype, data)
+ else:
+ data.data = sc.bins_like(prototype, data.data)
+ return data
+
+
+
+@overload
+def drop_variances_if_broadcast(
+ data: sc.Variable, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.Variable:
+ pass
+
+
+@overload
+def drop_variances_if_broadcast(
+ data: sc.DataArray, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.DataArray:
+ pass
+
+
+
+[docs]
+def drop_variances_if_broadcast(
+ data: sc.Variable | sc.DataArray, /, *, prototype: sc.DataArray | sc.Variable
+) -> sc.Variable | sc.DataArray:
+ """
+ Drop variances if the data is broadcasted.
+
+ Parameters
+ ----------
+ data:
+ The data to broadcast.
+ prototype:
+ Defines the new sizes (dims and shape).
+
+ Returns
+ -------
+ :
+ The data without variances if the data is broadcasted.
+ """
+ if _no_variance_broadcast(data, prototype=prototype):
+ return data
+ return sc.values(data)
+
+
+
+def _no_variance_broadcast(
+ data: sc.Variable | sc.DataArray, /, *, prototype: sc.Variable | sc.DataArray
+) -> bool:
+ if data.bins is not None:
+ raise ValueError("Cannot broadcast binned data.")
+ if data.variances is None:
+ return True
+ if prototype.bins is not None:
+ return False
+ sizes = prototype.sizes
+ return all(data.sizes.get(dim) == size for dim, size in sizes.items())
+
+
+def _fail(
+ data: sc.Variable | sc.DataArray, /, *, prototype: sc.Variable | sc.DataArray
+) -> sc.Variable | sc.DataArray:
+ # If there are variances, a subsequent broadcasting operation using Scipp will fail.
+ # Do nothing here.
+ return data
+
+
+broadcasters = {
+ UncertaintyBroadcastMode.drop: drop_variances_if_broadcast,
+ UncertaintyBroadcastMode.upper_bound: broadcast_with_upper_bound_variances,
+ UncertaintyBroadcastMode.fail: _fail,
+}
+
+
+
+[docs]
+def broadcast_uncertainties(
+ data: sc.Variable | sc.DataArray,
+ /,
+ *,
+ prototype: sc.DataArray | sc.Variable,
+ mode: UncertaintyBroadcastMode,
+) -> sc.Variable | sc.DataArray:
+ """Broadcast uncertainties using the specified mode.
+
+ Since Scipp raises an error when broadcasting data with variances, this function
+ provides an explicit way to handle variances during broadcast operations.
+
+ Parameters
+ ----------
+ data:
+ Data with uncertainties to broadcast.
+ prototype:
+ Prototype defining the new sizes (dims and shape, or binned data sizes).
+ mode:
+ Selected broadcast mode.
+
+ Returns
+ -------
+ :
+ Data with broadcast uncertainties.
+ """
+ return broadcasters[mode](data, prototype=prototype)
+
+
+
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+# ruff: noqa: E402, F401
+from functools import singledispatch
+from typing import Any, Protocol
+
+import ipywidgets as widgets
+
+from ..parameter import (
+ BooleanParameter,
+ FilenameParameter,
+ MultiFilenameParameter,
+ ParamWithOptions,
+ StringParameter,
+ Parameter,
+ ParamWithBounds,
+ BinEdgesParameter,
+ Vector2dParameter,
+ Vector3dParameter,
+)
+from ._config import default_layout, default_style
+
+from ._binedges_widget import BinEdgesWidget
+from ._linspace_widget import LinspaceWidget
+from ._vector_widget import VectorWidget
+from ._bounds_widget import BoundsWidget
+from ._switchable_widget import SwitchWidget
+from ._optional_widget import OptionalWidget
+
+
+
+[docs]
+class EssWidget(Protocol):
+ """Protocol for ESS widgets.
+
+ All widgets should have a `value` property that returns the value of the widget.
+ It can be composed from multiple widgets.
+ ```
+ """
+
+ @property
+ def value(self) -> Any: ...
+
+
+
+from collections.abc import Callable
+from functools import wraps
+
+
+
+[docs]
+def switchable_widget(
+ func: Callable[[Parameter], widgets.Widget],
+) -> Callable[[Parameter], widgets.Widget]:
+ """Wrap a widget in a switchable widget."""
+
+ @wraps(func)
+ def wrapper(param: Parameter) -> widgets.Widget:
+ widget = func(param)
+ if param.switchable:
+ return SwitchWidget(widget, name=param.name)
+ return widget
+
+ return wrapper
+
+
+
+
+[docs]
+def optional_widget(
+ func: Callable[[Parameter], widgets.Widget],
+) -> Callable[[Parameter], widgets.Widget]:
+ """Wrap a widget in a optional widget."""
+
+ @wraps(func)
+ def wrapper(param: Parameter) -> widgets.Widget:
+ widget = func(param)
+ if param.optional:
+ return OptionalWidget(widget, name=param.name)
+ return widget
+
+ return wrapper
+
+
+
+
+[docs]
+@switchable_widget
+@optional_widget # optional_widget should be applied first
+@singledispatch
+def create_parameter_widget(param: Parameter) -> widgets.Widget:
+ """Create a widget for a parameter depending on the ``param`` type.
+
+ If the type of the parameter is not supported, a text widget is returned.
+ """
+ return widgets.Text(
+ '', description=param.name, layout=default_layout, style=default_style
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(BooleanParameter)
+def boolean_parameter_widget(param: BooleanParameter):
+ name = param.name.split('.')[-1]
+ description = param.description
+ return widgets.Checkbox(
+ value=param.default,
+ description=name,
+ tooltip=description,
+ layout=default_layout,
+ style=default_style,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(StringParameter)
+def string_parameter_widget(param: StringParameter):
+ name = param.name
+ description = param.description
+ if param.switchable:
+ return widgets.Text(
+ description=name,
+ tooltip=description,
+ layout=default_layout,
+ style=default_style,
+ )
+ else:
+ return widgets.Text(
+ value=param.default,
+ description=name,
+ tooltip=description,
+ layout=default_layout,
+ style=default_style,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(FilenameParameter)
+def filename_parameter_widget(param: FilenameParameter):
+ return widgets.Text(
+ description=param.name,
+ layout=default_layout,
+ style=default_style,
+ value=param.default,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(MultiFilenameParameter)
+def multi_filename_parameter_widget(param: MultiFilenameParameter):
+ return widgets.Text(
+ description=param.name,
+ layout=default_layout,
+ style=default_style,
+ value=param.default,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(ParamWithOptions)
+def param_with_option_widget(param: ParamWithOptions):
+ return widgets.Dropdown(
+ description=param.name,
+ options=param.options,
+ layout=default_layout,
+ style=default_style,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(ParamWithBounds)
+def param_with_bounds_widget(param: ParamWithBounds):
+ return BoundsWidget()
+
+
+
+
+[docs]
+@create_parameter_widget.register(BinEdgesParameter)
+def bin_edges_parameter_widget(param: BinEdgesParameter):
+ return BinEdgesWidget(
+ name=param.name,
+ dim=param.dim,
+ start=param.start,
+ stop=param.stop,
+ nbins=param.nbins,
+ unit=param.unit,
+ log=param.log,
+ )
+
+
+
+
+[docs]
+@create_parameter_widget.register(Vector2dParameter)
+def vector_2d_parameter_widget(param: Vector2dParameter):
+ return VectorWidget(name=param.name, variable=param.default, components="xy")
+
+
+
+
+[docs]
+@create_parameter_widget.register(Vector3dParameter)
+def vector_3d_parameter_widget(param: Vector3dParameter):
+ return VectorWidget(name=param.name, variable=param.default, components="xyz")
+
+
+
+__all__ = [
+ 'BinEdgesWidget',
+ 'BoundsWidget',
+ 'EssWidget',
+ 'LinspaceWidget',
+ 'OptionalWidget',
+ 'SwitchWidget',
+ 'VectorWidget',
+ 'create_parameter_widget',
+]
+
+ 
+
+
+
+
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2022 Scipp contributors (https://github.com/scipp)
+# @author Simon Heybrock
+from __future__ import annotations
+
+import copy
+import functools
+import itertools
+import numbers
+import operator
+from collections.abc import (
+ Callable,
+ ItemsView,
+ Iterable,
+ Iterator,
+ KeysView,
+ Mapping,
+ MutableMapping,
+ Sequence,
+ ValuesView,
+)
+from functools import wraps
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Concatenate,
+ NoReturn,
+ ParamSpec,
+ TypeVar,
+ cast,
+ overload,
+)
+
+import numpy as np
+
+from .. import _binding
+from .cpp_classes import (
+ DataArray,
+ Dataset,
+ DimensionError,
+ GroupByDataArray,
+ GroupByDataset,
+ Unit,
+ Variable,
+)
+
+if TYPE_CHECKING:
+ # Avoid cyclic imports
+ from ..coords.graph import GraphDict
+ from ..typing import ScippIndex
+ from .bins import Bins
+
+
+_T = TypeVar("_T") # Any type
+_V = TypeVar("_V") # Value type of self
+_R = TypeVar("_R") # Return type of a callable
+_P = ParamSpec('_P')
+
+
+def _item_dims(item: Any) -> tuple[str, ...]:
+ return getattr(item, 'dims', ())
+
+
+def _is_binned(item: Any) -> bool:
+ from .bins import Bins
+
+ if isinstance(item, Bins):
+ return True
+ return getattr(item, 'bins', None) is not None
+
+
+def _summarize(item: Any) -> str:
+ if isinstance(item, DataGroup):
+ return f'{type(item).__name__}({len(item)}, {item.sizes})'
+ if hasattr(item, 'sizes'):
+ return f'{type(item).__name__}({item.sizes})'
+ return str(item)
+
+
+def _is_positional_index(key: Any) -> bool:
+ def is_int(x: object) -> bool:
+ return isinstance(x, numbers.Integral)
+
+ if is_int(key):
+ return True
+ if isinstance(key, slice):
+ if is_int(key.start) or is_int(key.stop) or is_int(key.step):
+ return True
+ if key.start is None and key.stop is None and key.step is None:
+ return True
+ return False
+
+
+def _is_list_index(key: Any) -> bool:
+ return isinstance(key, list | np.ndarray)
+
+
+class DataGroup(MutableMapping[str, _V]):
+ """
+ A dict-like group of data. Additionally provides dims and shape properties.
+
+ DataGroup acts like a Python dict but additionally supports Scipp functionality
+ such as positional- and label-based indexing and Scipp operations by mapping them
+ to the values in the dict. This may happen recursively to support tree-like data
+ structures.
+
+ .. versionadded:: 23.01.0
+ """
+
+ def __init__(
+ self, /, *args: Iterable[tuple[str, _V]] | Mapping[str, _V], **kwargs: _V
+ ) -> None:
+ self._items = dict(*args, **kwargs)
+ if not all(isinstance(k, str) for k in self._items.keys()):
+ raise ValueError("DataGroup keys must be strings.")
+
+ def __copy__(self) -> DataGroup[_V]:
+ return DataGroup(copy.copy(self._items))
+
+ def __len__(self) -> int:
+ """Return the number of items in the data group."""
+ return len(self._items)
+
+ def __iter__(self) -> Iterator[str]:
+ return iter(self._items)
+
+ def keys(self) -> KeysView[str]:
+ return self._items.keys()
+
+ def values(self) -> ValuesView[_V]:
+ return self._items.values()
+
+ def items(self) -> ItemsView[str, _V]:
+ return self._items.items()
+
+ @overload
+ def __getitem__(self, name: str) -> _V: ...
+
+ @overload
+ def __getitem__(self, name: ScippIndex) -> DataGroup[_V]: ...
+
+ def __getitem__(self, name: Any) -> Any:
+ """Return item of given name or index all items.
+
+ When ``name`` is a string, return the item of the given name. Otherwise, this
+ returns a new DataGroup, with items created by indexing the items in this
+ DataGroup. This may perform, e.g., Scipp's positional indexing, label-based
+ indexing, or advanced indexing on items that are scipp.Variable or
+ scipp.DataArray.
+
+ Label-based indexing is only possible when all items have a coordinate for the
+ indexed dimension.
+
+ Advanced indexing comprises integer-array indexing and boolean-variable
+ indexing. Unlike positional indexing, integer-array indexing works even when
+ the item shapes are inconsistent for the indexed dimensions, provided that all
+ items contain the maximal index in the integer array. Boolean-variable indexing
+ is only possible when the shape of all items is compatible with the boolean
+ variable.
+ """
+ from .bins import Bins
+
+ if isinstance(name, str):
+ return self._items[name]
+ if isinstance(name, tuple) and name == ():
+ return cast(DataGroup[Any], self).apply(operator.itemgetter(name))
+ if isinstance(name, Variable): # boolean indexing
+ return cast(DataGroup[Any], self).apply(operator.itemgetter(name))
+ if _is_positional_index(name) or _is_list_index(name):
+ if self.ndim != 1:
+ raise DimensionError(
+ "Slicing with implicit dimension label is only possible "
+ f"for 1-D objects. Got {self.sizes} with ndim={self.ndim}. Provide "
+ "an explicit dimension label, e.g., var['x', 0] instead of var[0]."
+ )
+ dim = self.dims[0]
+ index = name
+ else:
+ dim, index = name
+ return DataGroup(
+ {
+ key: var[dim, index] # type: ignore[index]
+ if (isinstance(var, Bins) or dim in _item_dims(var))
+ else var
+ for key, var in self.items()
+ }
+ )
+
+ def __setitem__(self, name: str, value: _V) -> None:
+ """Set self[key] to value."""
+ if isinstance(name, str):
+ self._items[name] = value
+ else:
+ raise TypeError('Keys must be strings')
+
+ def __delitem__(self, name: str) -> None:
+ """Delete self[key]."""
+ del self._items[name]
+
+ def __sizeof__(self) -> int:
+ return self.underlying_size()
+
+ def underlying_size(self) -> int:
+ # TODO Return the underlying size of all items in DataGroup
+ total_size = super.__sizeof__(self)
+ for item in self.values():
+ if isinstance(item, DataArray | Dataset | Variable | DataGroup):
+ total_size += item.underlying_size()
+ elif hasattr(item, 'nbytes'):
+ total_size += item.nbytes
+ else:
+ total_size += item.__sizeof__()
+
+ return total_size
+
+ @property
+ def dims(self) -> tuple[str, ...]:
+ """Union of dims of all items. Non-Scipp items are handled as dims=()."""
+ return tuple(self.sizes)
+
+ @property
+ def ndim(self) -> int:
+ """Number of dimensions, i.e., len(self.dims)."""
+ return len(self.dims)
+
+ @property
+ def shape(self) -> tuple[int | None, ...]:
+ """Union of shape of all items. Non-Scipp items are handled as shape=()."""
+ return tuple(self.sizes.values())
+
+ @property
+ def sizes(self) -> dict[str, int | None]:
+ """Dict combining dims and shape, i.e., mapping dim labels to their size."""
+ all_sizes: dict[str, set[int]] = {}
+ for x in self.values():
+ for dim, size in getattr(x, 'sizes', {}).items():
+ all_sizes.setdefault(dim, set()).add(size)
+ return {d: next(iter(s)) if len(s) == 1 else None for d, s in all_sizes.items()}
+
+ def _repr_html_(self) -> str:
+ from ..visualization.formatting_datagroup_html import datagroup_repr
+
+ return datagroup_repr(self)
+
+ def __repr__(self) -> str:
+ r = f'DataGroup(sizes={self.sizes}, keys=[\n'
+ for name, var in self.items():
+ r += f' {name}: {_summarize(var)},\n'
+ r += '])'
+ return r
+
+ def __str__(self) -> str:
+ return f'DataGroup(sizes={self.sizes}, keys={list(self.keys())})'
+
+ @property
+ def bins(self) -> DataGroup[Bins | None]:
+ # TODO Returning a regular DataGroup here may be wrong, since the `bins`
+ # property provides a different set of attrs and methods.
+ return self.apply(operator.attrgetter('bins'))
+
+ def apply(
+ self,
+ func: Callable[Concatenate[_V, _P], _R],
+ *args: _P.args,
+ **kwargs: _P.kwargs,
+ ) -> DataGroup[_R]:
+ """Call func on all values and return new DataGroup containing the results."""
+ return DataGroup({key: func(v, *args, **kwargs) for key, v in self.items()})
+
+ def _transform_dim(
+ self, func: str, *, dim: None | str | Iterable[str], **kwargs: Any
+ ) -> DataGroup[Any]:
+ """Transform items that depend on one or more dimensions given by `dim`."""
+ dims = (dim,) if isinstance(dim, str) else dim
+
+ def intersects(item: _V) -> bool:
+ item_dims = _item_dims(item)
+ if dims is None:
+ return item_dims != ()
+ return set(dims).intersection(item_dims) != set()
+
+ return DataGroup(
+ {
+ key: v
+ if not intersects(v)
+ else operator.methodcaller(func, dim, **kwargs)(v)
+ for key, v in self.items()
+ }
+ )
+
+ def _reduce(
+ self, method: str, dim: None | str | Sequence[str] = None, **kwargs: Any
+ ) -> DataGroup[Any]:
+ reduce_all = operator.methodcaller(method, **kwargs)
+
+ def _reduce_child(v: _V) -> Any:
+ if isinstance(v, GroupByDataArray | GroupByDataset):
+ child_dims: tuple[None | str | Sequence[str], ...] = (dim,)
+ else:
+ child_dims = _item_dims(v)
+ # Reduction operations on binned data implicitly reduce over bin content.
+ # Therefore, a purely dimension-based logic is not sufficient to determine
+ # if the item has to be reduced or not.
+ binned = _is_binned(v)
+ if child_dims == () and not binned:
+ return v
+ if dim is None:
+ return reduce_all(v)
+ if isinstance(dim, str):
+ dims_to_reduce = dim if dim in child_dims else ()
+ else:
+ dims_to_reduce = tuple(d for d in dim if d in child_dims)
+ if dims_to_reduce == () and binned:
+ return reduce_all(v)
+ return (
+ v
+ if dims_to_reduce == ()
+ else operator.methodcaller(method, dims_to_reduce, **kwargs)(v)
+ )
+
+ return DataGroup({key: _reduce_child(v) for key, v in self.items()})
+
+ def copy(self, deep: bool = True) -> DataGroup[_V]:
+ return copy.deepcopy(self) if deep else copy.copy(self)
+
+ def all(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('all', dim)
+
+ def any(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('any', dim)
+
+ def astype(self, type: Any, *, copy: bool = True) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('astype', type, copy=copy))
+
+ def bin(
+ self,
+ arg_dict: dict[str, int | Variable] | None = None,
+ /,
+ **kwargs: int | Variable,
+ ) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('bin', arg_dict, **kwargs))
+
+ @overload
+ def broadcast(
+ self,
+ *,
+ dims: Sequence[str],
+ shape: Sequence[int],
+ ) -> DataGroup[_V]: ...
+
+ @overload
+ def broadcast(
+ self,
+ *,
+ sizes: dict[str, int],
+ ) -> DataGroup[_V]: ...
+
+ def broadcast(
+ self,
+ *,
+ dims: Sequence[str] | None = None,
+ shape: Sequence[int] | None = None,
+ sizes: dict[str, int] | None = None,
+ ) -> DataGroup[_V]:
+ return self.apply(
+ operator.methodcaller('broadcast', dims=dims, shape=shape, sizes=sizes)
+ )
+
+ def ceil(self) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('ceil'))
+
+ def flatten(
+ self, dims: Sequence[str] | None = None, to: str | None = None
+ ) -> DataGroup[_V]:
+ return self._transform_dim('flatten', dim=dims, to=to)
+
+ def floor(self) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('floor'))
+
+ @overload
+ def fold(
+ self,
+ dim: str,
+ *,
+ dims: Sequence[str],
+ shape: Sequence[int],
+ ) -> DataGroup[_V]: ...
+
+ @overload
+ def fold(
+ self,
+ dim: str,
+ *,
+ sizes: dict[str, int],
+ ) -> DataGroup[_V]: ...
+
+ def fold(
+ self,
+ dim: str,
+ *,
+ dims: Sequence[str] | None = None,
+ shape: Sequence[int] | None = None,
+ sizes: dict[str, int] | None = None,
+ ) -> DataGroup[_V]:
+ return self._transform_dim('fold', dim=dim, dims=dims, shape=shape, sizes=sizes)
+
+ def group(self, /, *args: str | Variable) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('group', *args))
+
+ def groupby(
+ self, /, group: Variable | str, *, bins: Variable | None = None
+ ) -> DataGroup[GroupByDataArray | GroupByDataset]:
+ return self.apply(operator.methodcaller('groupby', group, bins=bins))
+
+ def hist(
+ self,
+ arg_dict: dict[str, int | Variable] | None = None,
+ /,
+ **kwargs: int | Variable,
+ ) -> DataGroup[DataArray | Dataset]:
+ return self.apply(operator.methodcaller('hist', arg_dict, **kwargs))
+
+ def max(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('max', dim)
+
+ def mean(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('mean', dim)
+
+ def median(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('median', dim)
+
+ def min(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('min', dim)
+
+ def nanhist(
+ self,
+ arg_dict: dict[str, int | Variable] | None = None,
+ /,
+ **kwargs: int | Variable,
+ ) -> DataGroup[DataArray]:
+ return self.apply(operator.methodcaller('nanhist', arg_dict, **kwargs))
+
+ def nanmax(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('nanmax', dim)
+
+ def nanmean(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('nanmean', dim)
+
+ def nanmedian(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('nanmedian', dim)
+
+ def nanmin(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('nanmin', dim)
+
+ def nansum(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('nansum', dim)
+
+ def nanstd(
+ self, dim: None | str | tuple[str, ...] = None, *, ddof: int
+ ) -> DataGroup[_V]:
+ return self._reduce('nanstd', dim, ddof=ddof)
+
+ def nanvar(
+ self, dim: None | str | tuple[str, ...] = None, *, ddof: int
+ ) -> DataGroup[_V]:
+ return self._reduce('nanvar', dim, ddof=ddof)
+
+ def rebin(
+ self,
+ arg_dict: dict[str, int | Variable] | None = None,
+ /,
+ **kwargs: int | Variable,
+ ) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('rebin', arg_dict, **kwargs))
+
+ def rename(
+ self, dims_dict: dict[str, str] | None = None, /, **names: str
+ ) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('rename', dims_dict, **names))
+
+ def rename_dims(
+ self, dims_dict: dict[str, str] | None = None, /, **names: str
+ ) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('rename_dims', dims_dict, **names))
+
+ def round(self) -> DataGroup[_V]:
+ return self.apply(operator.methodcaller('round'))
+
+ def squeeze(self, dim: str | Sequence[str] | None = None) -> DataGroup[_V]:
+ return self._reduce('squeeze', dim)
+
+ def std(
+ self, dim: None | str | tuple[str, ...] = None, *, ddof: int
+ ) -> DataGroup[_V]:
+ return self._reduce('std', dim, ddof=ddof)
+
+ def sum(self, dim: None | str | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._reduce('sum', dim)
+
+ def to(
+ self,
+ *,
+ unit: Unit | str | None = None,
+ dtype: Any | None = None,
+ copy: bool = True,
+ ) -> DataGroup[_V]:
+ return self.apply(
+ operator.methodcaller('to', unit=unit, dtype=dtype, copy=copy)
+ )
+
+ def transform_coords(
+ self,
+ targets: str | Iterable[str] | None = None,
+ /,
+ graph: GraphDict | None = None,
+ *,
+ rename_dims: bool = True,
+ keep_aliases: bool = True,
+ keep_intermediate: bool = True,
+ keep_inputs: bool = True,
+ quiet: bool = False,
+ **kwargs: Callable[..., Variable],
+ ) -> DataGroup[_V]:
+ return self.apply(
+ operator.methodcaller(
+ 'transform_coords',
+ targets,
+ graph=graph,
+ rename_dims=rename_dims,
+ keep_aliases=keep_aliases,
+ keep_intermediate=keep_intermediate,
+ keep_inputs=keep_inputs,
+ quiet=quiet,
+ **kwargs,
+ )
+ )
+
+ def transpose(self, dims: None | tuple[str, ...] = None) -> DataGroup[_V]:
+ return self._transform_dim('transpose', dim=dims)
+
+ def var(
+ self, dim: None | str | tuple[str, ...] = None, *, ddof: int
+ ) -> DataGroup[_V]:
+ return self._reduce('var', dim, ddof=ddof)
+
+ def plot(self, *args: Any, **kwargs: Any) -> Any:
+ import plopp
+
+ return plopp.plot(self, *args, **kwargs)
+
+ def __eq__( # type: ignore[override]
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise equal."""
+ return data_group_nary(operator.eq, self, other)
+
+ def __ne__( # type: ignore[override]
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise not-equal."""
+ return data_group_nary(operator.ne, self, other)
+
+ def __gt__(
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise greater-than."""
+ return data_group_nary(operator.gt, self, other)
+
+ def __ge__(
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise greater-equal."""
+ return data_group_nary(operator.ge, self, other)
+
+ def __lt__(
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise less-than."""
+ return data_group_nary(operator.lt, self, other)
+
+ def __le__(
+ self, other: DataGroup[object] | DataArray | Variable | float
+ ) -> DataGroup[_V | bool]:
+ """Item-wise less-equal."""
+ return data_group_nary(operator.le, self, other)
+
+ def __add__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``add`` item-by-item."""
+ return data_group_nary(operator.add, self, other)
+
+ def __sub__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``sub`` item-by-item."""
+ return data_group_nary(operator.sub, self, other)
+
+ def __mul__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``mul`` item-by-item."""
+ return data_group_nary(operator.mul, self, other)
+
+ def __truediv__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``truediv`` item-by-item."""
+ return data_group_nary(operator.truediv, self, other)
+
+ def __floordiv__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``floordiv`` item-by-item."""
+ return data_group_nary(operator.floordiv, self, other)
+
+ def __mod__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``mod`` item-by-item."""
+ return data_group_nary(operator.mod, self, other)
+
+ def __pow__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Apply ``pow`` item-by-item."""
+ return data_group_nary(operator.pow, self, other)
+
+ def __radd__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``add`` item-by-item."""
+ return data_group_nary(operator.add, other, self)
+
+ def __rsub__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``sub`` item-by-item."""
+ return data_group_nary(operator.sub, other, self)
+
+ def __rmul__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``mul`` item-by-item."""
+ return data_group_nary(operator.mul, other, self)
+
+ def __rtruediv__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``truediv`` item-by-item."""
+ return data_group_nary(operator.truediv, other, self)
+
+ def __rfloordiv__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``floordiv`` item-by-item."""
+ return data_group_nary(operator.floordiv, other, self)
+
+ def __rmod__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``mod`` item-by-item."""
+ return data_group_nary(operator.mod, other, self)
+
+ def __rpow__(self, other: DataArray | Variable | float) -> DataGroup[Any]:
+ """Apply ``pow`` item-by-item."""
+ return data_group_nary(operator.pow, other, self)
+
+ def __and__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Return the element-wise ``and`` of items."""
+ return data_group_nary(operator.and_, self, other)
+
+ def __or__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Return the element-wise ``or`` of items."""
+ return data_group_nary(operator.or_, self, other)
+
+ def __xor__(
+ self, other: DataGroup[Any] | DataArray | Variable | float
+ ) -> DataGroup[Any]:
+ """Return the element-wise ``xor`` of items."""
+ return data_group_nary(operator.xor, self, other)
+
+ def __invert__(self) -> DataGroup[Any]:
+ """Return the element-wise ``or`` of items."""
+ return self.apply(operator.invert) # type: ignore[arg-type]
+
+
+def data_group_nary(
+ func: Callable[..., _R], *args: Any, **kwargs: Any
+) -> DataGroup[_R]:
+ dgs = filter(
+ lambda x: isinstance(x, DataGroup), itertools.chain(args, kwargs.values())
+ )
+ keys = functools.reduce(operator.and_, [dg.keys() for dg in dgs])
+
+ def elem(x: Any, key: str) -> Any:
+ return x[key] if isinstance(x, DataGroup) else x
+
+ return DataGroup(
+ {
+ key: func(
+ *[elem(x, key) for x in args],
+ **{name: elem(x, key) for name, x in kwargs.items()},
+ )
+ for key in keys
+ }
+ )
+
+
+def apply_to_items(
+ func: Callable[..., _R], dgs: Iterable[DataGroup[Any]], *args: Any, **kwargs: Any
+) -> DataGroup[_R]:
+ keys = functools.reduce(operator.and_, [dg.keys() for dg in dgs])
+ return DataGroup(
+ {key: func([dg[key] for dg in dgs], *args, **kwargs) for key in keys}
+ )
+
+
+def data_group_overload(
+ func: Callable[Concatenate[_T, _P], _R],
+) -> Callable[..., _R | DataGroup[_R]]:
+ """Add an overload for DataGroup to a function.
+
+ If the first argument of the function is a data group,
+ then the decorated function is mapped over all items.
+ It is applied recursively for items that are themselves data groups.
+
+ Otherwise, the original function is applied directly.
+
+ Parameters
+ ----------
+ func:
+ Function to decorate.
+
+ Returns
+ -------
+ :
+ Decorated function.
+ """
+
+ # Do not assign '__annotations__' because that causes an error in Sphinx.
+ @wraps(func, assigned=('__module__', '__name__', '__qualname__', '__doc__'))
+ def impl(
+ data: _T | DataGroup[Any], *args: _P.args, **kwargs: _P.kwargs
+ ) -> _R | DataGroup[_R]:
+ if isinstance(data, DataGroup):
+ return data.apply(impl, *args, **kwargs) # type: ignore[arg-type]
+ return func(data, *args, **kwargs)
+
+ return impl
+
+
+# There are currently no in-place operations (__iadd__, etc.) because they require
+# a check if the operation would fail before doing it. As otherwise, a failure could
+# leave a partially modified data group behind. Dataset implements such a check, but
+# it is simpler than for DataGroup because the latter supports more data types.
+# So for now, we went with the simple solution and
+# not support in-place operations at all.
+#
+# Binding these functions dynamically has the added benefit that type checkers think
+# that the operations are not implemented.
+def _make_inplace_binary_op(name: str) -> Callable[..., NoReturn]:
+ def impl(
+ self: DataGroup[Any], other: DataGroup[Any] | DataArray | Variable | float
+ ) -> NoReturn:
+ raise TypeError(f'In-place operation i{name} is not supported by DataGroup.')
+
+ return impl
+
+
+for _name in ('add', 'sub', 'mul', 'truediv', 'floordiv', 'mod', 'pow'):
+ full_name = f'__i{_name}__'
+ _binding.bind_function_as_method(
+ cls=DataGroup, name=full_name, func=_make_inplace_binary_op(full_name)
+ )
+
+del _name, full_name
+
+ 
+
+
+
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+# @author Simon Heybrock
+"""
+Utilities for loading and working with NeXus transformations.
+
+Transformation chains in NeXus files can be non-local and can thus be challenging to
+work with. Additionally, values of transformations can be time-dependent, with each
+chain link potentially having a different time-dependent value. In practice the user is
+interested in the position and orientation of a component at a specific time or time
+range. This may involve evaluating the transformation chain at a specific time, or
+applying some heuristic to determine if the changes in the transformation value are
+significant or just noise. In combination, the above means that we need to remain
+flexible in how we handle transformations, preserving all necessary information from
+the source files. Therefore:
+
+1. :py:class:`Transform` is a dataclass representing a transformation. The raw `value`
+ dataset is preserved (instead of directly converting to, e.g., a rotation matrix) to
+ facilitate further processing such as computing the mean or variance.
+2. Loading a :py:class:`Group` will follow depends_on chains and store them as an
+ attribute of the depends_on field. This is done by :py:func:`parse_depends_on_chain`.
+3. :py:func:`compute_positions` computes component positions (and transformations). By
+ making this an explicit separate step, transformations can be applied to the
+ transformations stored by the depends_on field before doing so. We imagine that this
+ can be used to
+
+ - Evaluate the transformation at a specific time.
+ - Apply filters to remove noise, to avoid having to deal with very small time
+ intervals when processing data.
+
+By keeping the loaded transformations in a simple and modifiable format, we can
+furthermore manually update the transformations with information from other sources,
+such as streamed NXlog values received from a data acquisition system.
+"""
+
+from __future__ import annotations
+
+import warnings
+from dataclasses import dataclass, field, replace
+from typing import Literal
+
+import numpy as np
+import scipp as sc
+from scipp.scipy import interpolate
+
+from .base import Group, NexusStructureError, NXobject, base_definitions_dict
+from .field import DependsOn, Field
+
+
+class NXtransformations(NXobject):
+ """
+ Group of transformations.
+
+ Currently all transformations in the group are loaded. This may lead to redundant
+ loads as transformations are also loaded by following depends_on chains.
+ """
+
+
+class TransformationError(NexusStructureError):
+ pass
+
+
+@dataclass
+class Transform:
+ """In-memory component translation or rotation as described by NXtransformations."""
+
+ name: str
+ transformation_type: Literal['translation', 'rotation']
+ value: sc.Variable | sc.DataArray | sc.DataGroup
+ vector: sc.Variable
+ depends_on: DependsOn
+ offset: sc.Variable | None = None
+
+ @property
+ def sizes(self) -> dict[str, int]:
+ """Convenience property to access sizes of the value."""
+ return self.value.sizes
+
+ def __post_init__(self):
+ if self.transformation_type not in ['translation', 'rotation']:
+ raise TransformationError(
+ f"{self.transformation_type=} attribute at {self.name},"
+ " expected 'translation' or 'rotation'."
+ )
+
+ @staticmethod
+ def from_object(
+ obj: Field | Group, value: sc.Variable | sc.DataArray | sc.DataGroup
+ ) -> Transform:
+ depends_on = DependsOn(parent=obj.parent.name, value=obj.attrs['depends_on'])
+ return Transform(
+ name=obj.name,
+ transformation_type=obj.attrs.get('transformation_type'),
+ value=_parse_value(value),
+ vector=sc.vector(value=obj.attrs['vector']),
+ depends_on=depends_on,
+ offset=_parse_offset(obj),
+ )
+
+ def build(self) -> sc.Variable | sc.DataArray:
+ """Convert the raw transform into a rotation or translation matrix."""
+ t = self.value * self.vector
+ v = t if isinstance(t, sc.Variable) else t.data
+ if self.transformation_type == 'translation':
+ v = sc.spatial.translations(dims=v.dims, values=v.values, unit=v.unit)
+ elif self.transformation_type == 'rotation':
+ v = sc.spatial.rotations_from_rotvecs(v)
+ if isinstance(t, sc.Variable):
+ t = v
+ else:
+ t.data = v
+ if self.offset is None:
+ return t
+ if self.transformation_type == 'translation':
+ return t * self.offset.to(unit=t.unit, copy=False)
+ return t * self.offset
+
+
+def _parse_offset(obj: Field | Group) -> sc.Variable | None:
+ if (offset := obj.attrs.get('offset')) is None:
+ return None
+ if (offset_units := obj.attrs.get('offset_units')) is None:
+ raise TransformationError(
+ f"Found {offset=} but no corresponding 'offset_units' "
+ f"attribute at {obj.name}"
+ )
+ return sc.spatial.translation(value=offset, unit=offset_units)
+
+
+def _parse_value(
+ value: sc.Variable | sc.DataArray | sc.DataGroup,
+) -> sc.Variable | sc.DataArray | sc.DataGroup:
+ if isinstance(value, sc.DataGroup) and (
+ isinstance(value.get('value'), sc.DataArray)
+ ):
+ # Some NXlog groups are split into value, alarm, and connection_status
+ # sublogs. We only care about the value.
+ value = value['value']
+ return value
+
+
+def _interpolate_transform(transform, xnew):
+ # scipy can't interpolate with a single value
+ if transform.sizes["time"] == 1:
+ transform = sc.concat([transform, transform], dim="time")
+ return interpolate.interp1d(
+ transform, "time", kind="previous", fill_value="extrapolate"
+ )(xnew=xnew)
+
+
+def _smaller_unit(a, b):
+ if a.unit == b.unit:
+ return a.unit
+ ratio = sc.scalar(1.0, unit=a.unit).to(unit=b.unit)
+ if ratio.value < 1.0:
+ return a.unit
+ else:
+ return b.unit
+
+
+def combine_transformations(
+ chain: list[sc.DataArray | sc.Variable],
+) -> sc.DataArray | sc.Variable:
+ """
+ Take the product of a chain of transformations, handling potentially mismatching
+ time-dependence.
+
+ Time-dependent transformations are interpolated to a common time-coordinate.
+ """
+ if any((x.sizes.get('time') == 0) for x in chain):
+ warnings.warn(
+ UserWarning('depends_on chain contains empty time-series, '), stacklevel=2
+ )
+ # It is not clear what the dtype should be in this case. As transformations
+ # are commonly multiplied onto position vectors, we return an empty array of
+ # floats, which can be multiplied by Scipp's vector dtype.
+ return sc.DataArray(
+ sc.array(dims=['time'], values=[], dtype='float64', unit=''),
+ coords={'time': sc.datetimes(dims=['time'], values=[], unit='s')},
+ )
+ total_transform = None
+ for transform in chain:
+ if transform.dtype in (sc.DType.translation3, sc.DType.affine_transform3):
+ transform = transform.to(unit='m', copy=False)
+ if total_transform is None:
+ total_transform = transform
+ elif isinstance(total_transform, sc.DataArray) and isinstance(
+ transform, sc.DataArray
+ ):
+ unit = _smaller_unit(
+ transform.coords['time'], total_transform.coords['time']
+ )
+ total_transform.coords['time'] = total_transform.coords['time'].to(
+ unit=unit, copy=False
+ )
+ transform.coords['time'] = transform.coords['time'].to(
+ unit=unit, copy=False
+ )
+ time = sc.concat(
+ [total_transform.coords["time"], transform.coords["time"]], dim="time"
+ )
+ time = sc.datetimes(values=np.unique(time.values), dims=["time"], unit=unit)
+ total_transform = _interpolate_transform(
+ transform, time
+ ) * _interpolate_transform(total_transform, time)
+ else:
+ total_transform = transform * total_transform
+ if isinstance(total_transform, sc.DataArray):
+ time_dependent = [t for t in chain if isinstance(t, sc.DataArray)]
+ times = [da.coords['time'][0] for da in time_dependent]
+ latest_log_start = sc.reduce(times).max()
+ return total_transform['time', latest_log_start:].copy()
+ return sc.scalar(1) if total_transform is None else total_transform
+
+
+def maybe_transformation(
+ obj: Field | Group, value: sc.Variable | sc.DataArray | sc.DataGroup
+) -> sc.Variable | sc.DataArray | sc.DataGroup:
+ """
+ Return a loaded field, possibly modified if it is a transformation.
+
+ Transformations are usually stored in NXtransformations groups. However, identifying
+ transformation fields in this way requires inspecting the parent group, which
+ is cumbersome to implement. Furthermore, according to the NXdetector documentation
+ transformations are not necessarily placed inside NXtransformations.
+ Instead we use the presence of the attribute 'transformation_type' to identify
+ transformation fields.
+ """
+ if obj.attrs.get('transformation_type') is None:
+ return value
+ try:
+ return Transform.from_object(obj, value)
+ except KeyError as e:
+ warnings.warn(
+ UserWarning(f'Invalid transformation, missing attribute {e}'), stacklevel=2
+ )
+ return value
+
+
+@dataclass
+class TransformationChain(DependsOn):
+ """
+ Represents a chain of transformations references by a depends_on field.
+
+ Loading a group with a depends_on field will try to follow the chain and store the
+ transformations as an additional attribute of the in-memory representation of the
+ depends_on field.
+ """
+
+ transformations: sc.DataGroup = field(default_factory=sc.DataGroup)
+
+ def compute(self) -> sc.Variable | sc.DataArray:
+ depends_on = self
+ try:
+ chain = []
+ while (path := depends_on.absolute_path()) is not None:
+ chain.append(self.transformations[path])
+ depends_on = chain[-1].depends_on
+ transform = combine_transformations([t.build() for t in chain])
+ except KeyError as e:
+ warnings.warn(
+ UserWarning(f'depends_on chain references missing node:\n{e}'),
+ stacklevel=2,
+ )
+ else:
+ return transform
+
+
+def parse_depends_on_chain(
+ parent: Field | Group, depends_on: DependsOn
+) -> TransformationChain | None:
+ """Follow a depends_on chain and return the transformations."""
+ chain = TransformationChain(depends_on.parent, depends_on.value)
+ depends_on = depends_on.value
+ try:
+ while depends_on != '.':
+ transform = parent[depends_on]
+ parent = transform.parent
+ depends_on = transform.attrs['depends_on']
+ chain.transformations[transform.name] = transform[()]
+ except KeyError as e:
+ warnings.warn(
+ UserWarning(f'depends_on chain references missing node {e}'), stacklevel=2
+ )
+ return None
+ return chain
+
+
+def compute_positions(
+ dg: sc.DataGroup,
+ *,
+ store_position: str = 'position',
+ store_transform: str | None = None,
+ transformations: sc.DataGroup | None = None,
+) -> sc.DataGroup:
+ """
+ Recursively compute positions from depends_on attributes as well as the
+ [xyz]_pixel_offset fields of NXdetector groups.
+
+ This function does not operate directly on a NeXus file but on the result of
+ loading a NeXus file or sub-group into a scipp.DataGroup. NeXus puts no
+ limitations on the structure of the depends_on chains, i.e., they may reference
+ parent groups. If this is the case, a call to this function will fail if only the
+ subgroup is passed as input.
+
+ Note that this does not consider "legacy" ways of storing positions. In particular,
+ ``NXmonitor.distance``, ``NXdetector.distance``, ``NXdetector.polar_angle``, and
+ ``NXdetector.azimuthal_angle`` are ignored.
+
+ Note that transformation chains may be time-dependent. In this case it will not
+ be applied to the pixel offsets, since the result may consume too much memory and
+ the shape is in general incompatible with the shape of the data. Use the
+ ``store_transform`` argument to store the resolved transformation chain in this
+ case.
+
+ If a transformation chain has an invalid 'depends_on' value, e.g., a path beyond
+ the root data group, then the chain is ignored and no position is computed. This
+ does not affect other chains.
+
+ Parameters
+ ----------
+ dg:
+ Data group with depends_on entry points into transformation chains.
+ store_position:
+ Name used to store result of resolving each depends_on chain.
+ store_transform:
+ If not None, store the resolved transformation chain in this field.
+ transformations:
+ Optional data group containing transformation chains. If not provided, the
+ transformations are looked up in the chains stored within the depends_on field.
+
+ Returns
+ -------
+ :
+ New data group with added positions.
+ """
+ return _with_positions(
+ dg,
+ store_position=store_position,
+ store_transform=store_transform,
+ transformations=transformations,
+ )
+
+
+def zip_pixel_offsets(x: dict[str, sc.Variable], /) -> sc.Variable:
+ """
+ Zip the x_pixel_offset, y_pixel_offset, and z_pixel_offset fields into a vector.
+
+ These fields originate from NXdetector groups. All but x_pixel_offset are optional,
+ e.g., for 2D detectors. Zero values for missing fields are assumed.
+
+ Parameters
+ ----------
+ mapping:
+ Mapping (typically a data group, or data array coords) containing
+ x_pixel_offset, y_pixel_offset, and z_pixel_offset.
+
+ Returns
+ -------
+ :
+ Vectors with pixel offsets.
+
+ See Also
+ --------
+ compute_positions
+ """
+ zero = sc.scalar(0.0, unit=x['x_pixel_offset'].unit)
+ return sc.spatial.as_vectors(
+ x['x_pixel_offset'],
+ x.get('y_pixel_offset', zero),
+ x.get('z_pixel_offset', zero),
+ )
+
+
+def _with_positions(
+ dg: sc.DataGroup,
+ *,
+ store_position: str,
+ store_transform: str | None = None,
+ transformations: sc.DataGroup | None = None,
+) -> sc.DataGroup:
+ out = sc.DataGroup()
+ if (chain := dg.get('depends_on')) is not None:
+ if not isinstance(chain, TransformationChain):
+ chain = TransformationChain(chain.parent, chain.value)
+ if transformations is not None:
+ chain = replace(chain, transformations=transformations)
+ if (transform := chain.compute()) is not None:
+ out[store_position] = transform * sc.vector([0, 0, 0], unit='m')
+ if store_transform is not None:
+ out[store_transform] = transform
+ for name, value in dg.items():
+ if isinstance(value, sc.DataGroup):
+ value = _with_positions(
+ value,
+ store_position=store_position,
+ store_transform=store_transform,
+ transformations=transformations,
+ )
+ elif (
+ isinstance(value, sc.DataArray)
+ and 'x_pixel_offset' in value.coords
+ # Transform can be time-dependent, do not apply it to offsets since
+ # result can be massive and is in general not compatible with the shape
+ # of the data.
+ and (transform is not None and transform.dims == ())
+ ):
+ offset = zip_pixel_offsets(value.coords).to(unit='m', copy=False)
+ value = value.assign_coords({store_position: transform * offset})
+ out[name] = value
+ return out
+
+
+base_definitions_dict['NXtransformations'] = NXtransformations
+
+
+
+
+
+"""
+The typing module: Support for gradual typing as defined by PEP 484.
+
+At large scale, the structure of the module is following:
+* Imports and exports, all public names should be explicitly added to __all__.
+* Internal helper functions: these should never be used in code outside this module.
+* _SpecialForm and its instances (special forms):
+ Any, NoReturn, ClassVar, Union, Optional, Concatenate
+* Classes whose instances can be type arguments in addition to types:
+ ForwardRef, TypeVar and ParamSpec
+* The core of internal generics API: _GenericAlias and _VariadicGenericAlias, the latter is
+ currently only used by Tuple and Callable. All subscripted types like X[int], Union[int, str],
+ etc., are instances of either of these classes.
+* The public counterpart of the generics API consists of two classes: Generic and Protocol.
+* Public helper functions: get_type_hints, overload, cast, no_type_check,
+ no_type_check_decorator.
+* Generic aliases for collections.abc ABCs and few additional protocols.
+* Special types: NewType, NamedTuple, TypedDict.
+* Wrapper submodules for re and io related types.
+"""
+
+from abc import abstractmethod, ABCMeta
+import collections
+import collections.abc
+import contextlib
+import functools
+import operator
+import re as stdlib_re # Avoid confusion with the re we export.
+import sys
+import types
+from types import WrapperDescriptorType, MethodWrapperType, MethodDescriptorType, GenericAlias
+
+# Please keep __all__ alphabetized within each category.
+__all__ = [
+ # Super-special typing primitives.
+ 'Annotated',
+ 'Any',
+ 'Callable',
+ 'ClassVar',
+ 'Concatenate',
+ 'Final',
+ 'ForwardRef',
+ 'Generic',
+ 'Literal',
+ 'Optional',
+ 'ParamSpec',
+ 'Protocol',
+ 'Tuple',
+ 'Type',
+ 'TypeVar',
+ 'Union',
+
+ # ABCs (from collections.abc).
+ 'AbstractSet', # collections.abc.Set.
+ 'ByteString',
+ 'Container',
+ 'ContextManager',
+ 'Hashable',
+ 'ItemsView',
+ 'Iterable',
+ 'Iterator',
+ 'KeysView',
+ 'Mapping',
+ 'MappingView',
+ 'MutableMapping',
+ 'MutableSequence',
+ 'MutableSet',
+ 'Sequence',
+ 'Sized',
+ 'ValuesView',
+ 'Awaitable',
+ 'AsyncIterator',
+ 'AsyncIterable',
+ 'Coroutine',
+ 'Collection',
+ 'AsyncGenerator',
+ 'AsyncContextManager',
+
+ # Structural checks, a.k.a. protocols.
+ 'Reversible',
+ 'SupportsAbs',
+ 'SupportsBytes',
+ 'SupportsComplex',
+ 'SupportsFloat',
+ 'SupportsIndex',
+ 'SupportsInt',
+ 'SupportsRound',
+
+ # Concrete collection types.
+ 'ChainMap',
+ 'Counter',
+ 'Deque',
+ 'Dict',
+ 'DefaultDict',
+ 'List',
+ 'OrderedDict',
+ 'Set',
+ 'FrozenSet',
+ 'NamedTuple', # Not really a type.
+ 'TypedDict', # Not really a type.
+ 'Generator',
+
+ # Other concrete types.
+ 'BinaryIO',
+ 'IO',
+ 'Match',
+ 'Pattern',
+ 'TextIO',
+
+ # One-off things.
+ 'AnyStr',
+ 'cast',
+ 'final',
+ 'get_args',
+ 'get_origin',
+ 'get_type_hints',
+ 'is_typeddict',
+ 'NewType',
+ 'no_type_check',
+ 'no_type_check_decorator',
+ 'NoReturn',
+ 'overload',
+ 'ParamSpecArgs',
+ 'ParamSpecKwargs',
+ 'runtime_checkable',
+ 'Text',
+ 'TYPE_CHECKING',
+ 'TypeAlias',
+ 'TypeGuard',
+]
+
+# The pseudo-submodules 're' and 'io' are part of the public
+# namespace, but excluded from __all__ because they might stomp on
+# legitimate imports of those modules.
+
+
+def _type_convert(arg, module=None, *, allow_special_forms=False):
+ """For converting None to type(None), and strings to ForwardRef."""
+ if arg is None:
+ return type(None)
+ if isinstance(arg, str):
+ return ForwardRef(arg, module=module, is_class=allow_special_forms)
+ return arg
+
+
+def _type_check(arg, msg, is_argument=True, module=None, *, allow_special_forms=False):
+ """Check that the argument is a type, and return it (internal helper).
+
+ As a special case, accept None and return type(None) instead. Also wrap strings
+ into ForwardRef instances. Consider several corner cases, for example plain
+ special forms like Union are not valid, while Union[int, str] is OK, etc.
+ The msg argument is a human-readable error message, e.g::
+
+ "Union[arg, ...]: arg should be a type."
+
+ We append the repr() of the actual value (truncated to 100 chars).
+ """
+ invalid_generic_forms = (Generic, Protocol)
+ if not allow_special_forms:
+ invalid_generic_forms += (ClassVar,)
+ if is_argument:
+ invalid_generic_forms += (Final,)
+
+ arg = _type_convert(arg, module=module, allow_special_forms=allow_special_forms)
+ if (isinstance(arg, _GenericAlias) and
+ arg.__origin__ in invalid_generic_forms):
+ raise TypeError(f"{arg} is not valid as type argument")
+ if arg in (Any, NoReturn, Final, TypeAlias):
+ return arg
+ if isinstance(arg, _SpecialForm) or arg in (Generic, Protocol):
+ raise TypeError(f"Plain {arg} is not valid as type argument")
+ if isinstance(arg, (type, TypeVar, ForwardRef, types.UnionType, ParamSpec,
+ ParamSpecArgs, ParamSpecKwargs)):
+ return arg
+ if not callable(arg):
+ raise TypeError(f"{msg} Got {arg!r:.100}.")
+ return arg
+
+
+def _is_param_expr(arg):
+ return arg is ... or isinstance(arg,
+ (tuple, list, ParamSpec, _ConcatenateGenericAlias))
+
+
+def _type_repr(obj):
+ """Return the repr() of an object, special-casing types (internal helper).
+
+ If obj is a type, we return a shorter version than the default
+ type.__repr__, based on the module and qualified name, which is
+ typically enough to uniquely identify a type. For everything
+ else, we fall back on repr(obj).
+ """
+ if isinstance(obj, types.GenericAlias):
+ return repr(obj)
+ if isinstance(obj, type):
+ if obj.__module__ == 'builtins':
+ return obj.__qualname__
+ return f'{obj.__module__}.{obj.__qualname__}'
+ if obj is ...:
+ return('...')
+ if isinstance(obj, types.FunctionType):
+ return obj.__name__
+ return repr(obj)
+
+
+def _collect_type_vars(types_, typevar_types=None):
+ """Collect all type variable contained
+ in types in order of first appearance (lexicographic order). For example::
+
+ _collect_type_vars((T, List[S, T])) == (T, S)
+ """
+ if typevar_types is None:
+ typevar_types = TypeVar
+ tvars = []
+ for t in types_:
+ if isinstance(t, typevar_types) and t not in tvars:
+ tvars.append(t)
+ if isinstance(t, (_GenericAlias, GenericAlias, types.UnionType)):
+ tvars.extend([t for t in t.__parameters__ if t not in tvars])
+ return tuple(tvars)
+
+
+def _check_generic(cls, parameters, elen):
+ """Check correct count for parameters of a generic cls (internal helper).
+ This gives a nice error message in case of count mismatch.
+ """
+ if not elen:
+ raise TypeError(f"{cls} is not a generic class")
+ alen = len(parameters)
+ if alen != elen:
+ raise TypeError(f"Too {'many' if alen > elen else 'few'} arguments for {cls};"
+ f" actual {alen}, expected {elen}")
+
+def _prepare_paramspec_params(cls, params):
+ """Prepares the parameters for a Generic containing ParamSpec
+ variables (internal helper).
+ """
+ # Special case where Z[[int, str, bool]] == Z[int, str, bool] in PEP 612.
+ if (len(cls.__parameters__) == 1
+ and params and not _is_param_expr(params[0])):
+ assert isinstance(cls.__parameters__[0], ParamSpec)
+ return (params,)
+ else:
+ _check_generic(cls, params, len(cls.__parameters__))
+ _params = []
+ # Convert lists to tuples to help other libraries cache the results.
+ for p, tvar in zip(params, cls.__parameters__):
+ if isinstance(tvar, ParamSpec) and isinstance(p, list):
+ p = tuple(p)
+ _params.append(p)
+ return tuple(_params)
+
+def _deduplicate(params):
+ # Weed out strict duplicates, preserving the first of each occurrence.
+ all_params = set(params)
+ if len(all_params) < len(params):
+ new_params = []
+ for t in params:
+ if t in all_params:
+ new_params.append(t)
+ all_params.remove(t)
+ params = new_params
+ assert not all_params, all_params
+ return params
+
+
+def _remove_dups_flatten(parameters):
+ """An internal helper for Union creation and substitution: flatten Unions
+ among parameters, then remove duplicates.
+ """
+ # Flatten out Union[Union[...], ...].
+ params = []
+ for p in parameters:
+ if isinstance(p, (_UnionGenericAlias, types.UnionType)):
+ params.extend(p.__args__)
+ elif isinstance(p, tuple) and len(p) > 0 and p[0] is Union:
+ params.extend(p[1:])
+ else:
+ params.append(p)
+
+ return tuple(_deduplicate(params))
+
+
+def _flatten_literal_params(parameters):
+ """An internal helper for Literal creation: flatten Literals among parameters"""
+ params = []
+ for p in parameters:
+ if isinstance(p, _LiteralGenericAlias):
+ params.extend(p.__args__)
+ else:
+ params.append(p)
+ return tuple(params)
+
+
+_cleanups = []
+
+
+def _tp_cache(func=None, /, *, typed=False):
+ """Internal wrapper caching __getitem__ of generic types with a fallback to
+ original function for non-hashable arguments.
+ """
+ def decorator(func):
+ cached = functools.lru_cache(typed=typed)(func)
+ _cleanups.append(cached.cache_clear)
+
+ @functools.wraps(func)
+ def inner(*args, **kwds):
+ try:
+ return cached(*args, **kwds)
+ except TypeError:
+ pass # All real errors (not unhashable args) are raised below.
+ return func(*args, **kwds)
+ return inner
+
+ if func is not None:
+ return decorator(func)
+
+ return decorator
+
+def _eval_type(t, globalns, localns, recursive_guard=frozenset()):
+ """Evaluate all forward references in the given type t.
+ For use of globalns and localns see the docstring for get_type_hints().
+ recursive_guard is used to prevent infinite recursion with a recursive
+ ForwardRef.
+ """
+ if isinstance(t, ForwardRef):
+ return t._evaluate(globalns, localns, recursive_guard)
+ if isinstance(t, (_GenericAlias, GenericAlias, types.UnionType)):
+ ev_args = tuple(_eval_type(a, globalns, localns, recursive_guard) for a in t.__args__)
+ if ev_args == t.__args__:
+ return t
+ if isinstance(t, GenericAlias):
+ return GenericAlias(t.__origin__, ev_args)
+ if isinstance(t, types.UnionType):
+ return functools.reduce(operator.or_, ev_args)
+ else:
+ return t.copy_with(ev_args)
+ return t
+
+
+class _Final:
+ """Mixin to prohibit subclassing"""
+
+ __slots__ = ('__weakref__',)
+
+ def __init_subclass__(self, /, *args, **kwds):
+ if '_root' not in kwds:
+ raise TypeError("Cannot subclass special typing classes")
+
+class _Immutable:
+ """Mixin to indicate that object should not be copied."""
+ __slots__ = ()
+
+ def __copy__(self):
+ return self
+
+ def __deepcopy__(self, memo):
+ return self
+
+
+# Internal indicator of special typing constructs.
+# See __doc__ instance attribute for specific docs.
+class _SpecialForm(_Final, _root=True):
+ __slots__ = ('_name', '__doc__', '_getitem')
+
+ def __init__(self, getitem):
+ self._getitem = getitem
+ self._name = getitem.__name__
+ self.__doc__ = getitem.__doc__
+
+ def __getattr__(self, item):
+ if item in {'__name__', '__qualname__'}:
+ return self._name
+
+ raise AttributeError(item)
+
+ def __mro_entries__(self, bases):
+ raise TypeError(f"Cannot subclass {self!r}")
+
+ def __repr__(self):
+ return 'typing.' + self._name
+
+ def __reduce__(self):
+ return self._name
+
+ def __call__(self, *args, **kwds):
+ raise TypeError(f"Cannot instantiate {self!r}")
+
+ def __or__(self, other):
+ return Union[self, other]
+
+ def __ror__(self, other):
+ return Union[other, self]
+
+ def __instancecheck__(self, obj):
+ raise TypeError(f"{self} cannot be used with isinstance()")
+
+ def __subclasscheck__(self, cls):
+ raise TypeError(f"{self} cannot be used with issubclass()")
+
+ @_tp_cache
+ def __getitem__(self, parameters):
+ return self._getitem(self, parameters)
+
+
+class _LiteralSpecialForm(_SpecialForm, _root=True):
+ def __getitem__(self, parameters):
+ if not isinstance(parameters, tuple):
+ parameters = (parameters,)
+ return self._getitem(self, *parameters)
+
+
+@_SpecialForm
+def Any(self, parameters):
+ """Special type indicating an unconstrained type.
+
+ - Any is compatible with every type.
+ - Any assumed to have all methods.
+ - All values assumed to be instances of Any.
+
+ Note that all the above statements are true from the point of view of
+ static type checkers. At runtime, Any should not be used with instance
+ or class checks.
+ """
+ raise TypeError(f"{self} is not subscriptable")
+
+@_SpecialForm
+def NoReturn(self, parameters):
+ """Special type indicating functions that never return.
+ Example::
+
+ from typing import NoReturn
+
+ def stop() -> NoReturn:
+ raise Exception('no way')
+
+ This type is invalid in other positions, e.g., ``List[NoReturn]``
+ will fail in static type checkers.
+ """
+ raise TypeError(f"{self} is not subscriptable")
+
+@_SpecialForm
+def ClassVar(self, parameters):
+ """Special type construct to mark class variables.
+
+ An annotation wrapped in ClassVar indicates that a given
+ attribute is intended to be used as a class variable and
+ should not be set on instances of that class. Usage::
+
+ class Starship:
+ stats: ClassVar[Dict[str, int]] = {} # class variable
+ damage: int = 10 # instance variable
+
+ ClassVar accepts only types and cannot be further subscribed.
+
+ Note that ClassVar is not a class itself, and should not
+ be used with isinstance() or issubclass().
+ """
+ item = _type_check(parameters, f'{self} accepts only single type.')
+ return _GenericAlias(self, (item,))
+
+@_SpecialForm
+def Final(self, parameters):
+ """Special typing construct to indicate final names to type checkers.
+
+ A final name cannot be re-assigned or overridden in a subclass.
+ For example:
+
+ MAX_SIZE: Final = 9000
+ MAX_SIZE += 1 # Error reported by type checker
+
+ class Connection:
+ TIMEOUT: Final[int] = 10
+
+ class FastConnector(Connection):
+ TIMEOUT = 1 # Error reported by type checker
+
+ There is no runtime checking of these properties.
+ """
+ item = _type_check(parameters, f'{self} accepts only single type.')
+ return _GenericAlias(self, (item,))
+
+@_SpecialForm
+def Union(self, parameters):
+ """Union type; Union[X, Y] means either X or Y.
+
+ To define a union, use e.g. Union[int, str]. Details:
+ - The arguments must be types and there must be at least one.
+ - None as an argument is a special case and is replaced by
+ type(None).
+ - Unions of unions are flattened, e.g.::
+
+ Union[Union[int, str], float] == Union[int, str, float]
+
+ - Unions of a single argument vanish, e.g.::
+
+ Union[int] == int # The constructor actually returns int
+
+ - Redundant arguments are skipped, e.g.::
+
+ Union[int, str, int] == Union[int, str]
+
+ - When comparing unions, the argument order is ignored, e.g.::
+
+ Union[int, str] == Union[str, int]
+
+ - You cannot subclass or instantiate a union.
+ - You can use Optional[X] as a shorthand for Union[X, None].
+ """
+ if parameters == ():
+ raise TypeError("Cannot take a Union of no types.")
+ if not isinstance(parameters, tuple):
+ parameters = (parameters,)
+ msg = "Union[arg, ...]: each arg must be a type."
+ parameters = tuple(_type_check(p, msg) for p in parameters)
+ parameters = _remove_dups_flatten(parameters)
+ if len(parameters) == 1:
+ return parameters[0]
+ if len(parameters) == 2 and type(None) in parameters:
+ return _UnionGenericAlias(self, parameters, name="Optional")
+ return _UnionGenericAlias(self, parameters)
+
+@_SpecialForm
+def Optional(self, parameters):
+ """Optional type.
+
+ Optional[X] is equivalent to Union[X, None].
+ """
+ arg = _type_check(parameters, f"{self} requires a single type.")
+ return Union[arg, type(None)]
+
+@_LiteralSpecialForm
+@_tp_cache(typed=True)
+def Literal(self, *parameters):
+ """Special typing form to define literal types (a.k.a. value types).
+
+ This form can be used to indicate to type checkers that the corresponding
+ variable or function parameter has a value equivalent to the provided
+ literal (or one of several literals):
+
+ def validate_simple(data: Any) -> Literal[True]: # always returns True
+ ...
+
+ MODE = Literal['r', 'rb', 'w', 'wb']
+ def open_helper(file: str, mode: MODE) -> str:
+ ...
+
+ open_helper('/some/path', 'r') # Passes type check
+ open_helper('/other/path', 'typo') # Error in type checker
+
+ Literal[...] cannot be subclassed. At runtime, an arbitrary value
+ is allowed as type argument to Literal[...], but type checkers may
+ impose restrictions.
+ """
+ # There is no '_type_check' call because arguments to Literal[...] are
+ # values, not types.
+ parameters = _flatten_literal_params(parameters)
+
+ try:
+ parameters = tuple(p for p, _ in _deduplicate(list(_value_and_type_iter(parameters))))
+ except TypeError: # unhashable parameters
+ pass
+
+ return _LiteralGenericAlias(self, parameters)
+
+
+@_SpecialForm
+def TypeAlias(self, parameters):
+ """Special marker indicating that an assignment should
+ be recognized as a proper type alias definition by type
+ checkers.
+
+ For example::
+
+ Predicate: TypeAlias = Callable[..., bool]
+
+ It's invalid when used anywhere except as in the example above.
+ """
+ raise TypeError(f"{self} is not subscriptable")
+
+
+@_SpecialForm
+def Concatenate(self, parameters):
+ """Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a
+ higher order function which adds, removes or transforms parameters of a
+ callable.
+
+ For example::
+
+ Callable[Concatenate[int, P], int]
+
+ See PEP 612 for detailed information.
+ """
+ if parameters == ():
+ raise TypeError("Cannot take a Concatenate of no types.")
+ if not isinstance(parameters, tuple):
+ parameters = (parameters,)
+ if not isinstance(parameters[-1], ParamSpec):
+ raise TypeError("The last parameter to Concatenate should be a "
+ "ParamSpec variable.")
+ msg = "Concatenate[arg, ...]: each arg must be a type."
+ parameters = (*(_type_check(p, msg) for p in parameters[:-1]), parameters[-1])
+ return _ConcatenateGenericAlias(self, parameters,
+ _typevar_types=(TypeVar, ParamSpec),
+ _paramspec_tvars=True)
+
+
+@_SpecialForm
+def TypeGuard(self, parameters):
+ """Special typing form used to annotate the return type of a user-defined
+ type guard function. ``TypeGuard`` only accepts a single type argument.
+ At runtime, functions marked this way should return a boolean.
+
+ ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static
+ type checkers to determine a more precise type of an expression within a
+ program's code flow. Usually type narrowing is done by analyzing
+ conditional code flow and applying the narrowing to a block of code. The
+ conditional expression here is sometimes referred to as a "type guard".
+
+ Sometimes it would be convenient to use a user-defined boolean function
+ as a type guard. Such a function should use ``TypeGuard[...]`` as its
+ return type to alert static type checkers to this intention.
+
+ Using ``-> TypeGuard`` tells the static type checker that for a given
+ function:
+
+ 1. The return value is a boolean.
+ 2. If the return value is ``True``, the type of its argument
+ is the type inside ``TypeGuard``.
+
+ For example::
+
+ def is_str(val: Union[str, float]):
+ # "isinstance" type guard
+ if isinstance(val, str):
+ # Type of ``val`` is narrowed to ``str``
+ ...
+ else:
+ # Else, type of ``val`` is narrowed to ``float``.
+ ...
+
+ Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower
+ form of ``TypeA`` (it can even be a wider form) and this may lead to
+ type-unsafe results. The main reason is to allow for things like
+ narrowing ``List[object]`` to ``List[str]`` even though the latter is not
+ a subtype of the former, since ``List`` is invariant. The responsibility of
+ writing type-safe type guards is left to the user.
+
+ ``TypeGuard`` also works with type variables. For more information, see
+ PEP 647 (User-Defined Type Guards).
+ """
+ item = _type_check(parameters, f'{self} accepts only single type.')
+ return _GenericAlias(self, (item,))
+
+
+class ForwardRef(_Final, _root=True):
+ """Internal wrapper to hold a forward reference."""
+
+ __slots__ = ('__forward_arg__', '__forward_code__',
+ '__forward_evaluated__', '__forward_value__',
+ '__forward_is_argument__', '__forward_is_class__',
+ '__forward_module__')
+
+ def __init__(self, arg, is_argument=True, module=None, *, is_class=False):
+ if not isinstance(arg, str):
+ raise TypeError(f"Forward reference must be a string -- got {arg!r}")
+ try:
+ code = compile(arg, '<string>', 'eval')
+ except SyntaxError:
+ raise SyntaxError(f"Forward reference must be an expression -- got {arg!r}")
+ self.__forward_arg__ = arg
+ self.__forward_code__ = code
+ self.__forward_evaluated__ = False
+ self.__forward_value__ = None
+ self.__forward_is_argument__ = is_argument
+ self.__forward_is_class__ = is_class
+ self.__forward_module__ = module
+
+ def _evaluate(self, globalns, localns, recursive_guard):
+ if self.__forward_arg__ in recursive_guard:
+ return self
+ if not self.__forward_evaluated__ or localns is not globalns:
+ if globalns is None and localns is None:
+ globalns = localns = {}
+ elif globalns is None:
+ globalns = localns
+ elif localns is None:
+ localns = globalns
+ if self.__forward_module__ is not None:
+ globalns = getattr(
+ sys.modules.get(self.__forward_module__, None), '__dict__', globalns
+ )
+ type_ = _type_check(
+ eval(self.__forward_code__, globalns, localns),
+ "Forward references must evaluate to types.",
+ is_argument=self.__forward_is_argument__,
+ allow_special_forms=self.__forward_is_class__,
+ )
+ self.__forward_value__ = _eval_type(
+ type_, globalns, localns, recursive_guard | {self.__forward_arg__}
+ )
+ self.__forward_evaluated__ = True
+ return self.__forward_value__
+
+ def __eq__(self, other):
+ if not isinstance(other, ForwardRef):
+ return NotImplemented
+ if self.__forward_evaluated__ and other.__forward_evaluated__:
+ return (self.__forward_arg__ == other.__forward_arg__ and
+ self.__forward_value__ == other.__forward_value__)
+ return (self.__forward_arg__ == other.__forward_arg__ and
+ self.__forward_module__ == other.__forward_module__)
+
+ def __hash__(self):
+ return hash((self.__forward_arg__, self.__forward_module__))
+
+ def __repr__(self):
+ return f'ForwardRef({self.__forward_arg__!r})'
+
+class _TypeVarLike:
+ """Mixin for TypeVar-like types (TypeVar and ParamSpec)."""
+ def __init__(self, bound, covariant, contravariant):
+ """Used to setup TypeVars and ParamSpec's bound, covariant and
+ contravariant attributes.
+ """
+ if covariant and contravariant:
+ raise ValueError("Bivariant types are not supported.")
+ self.__covariant__ = bool(covariant)
+ self.__contravariant__ = bool(contravariant)
+ if bound:
+ self.__bound__ = _type_check(bound, "Bound must be a type.")
+ else:
+ self.__bound__ = None
+
+ def __or__(self, right):
+ return Union[self, right]
+
+ def __ror__(self, left):
+ return Union[left, self]
+
+ def __repr__(self):
+ if self.__covariant__:
+ prefix = '+'
+ elif self.__contravariant__:
+ prefix = '-'
+ else:
+ prefix = '~'
+ return prefix + self.__name__
+
+ def __reduce__(self):
+ return self.__name__
+
+
+class TypeVar( _Final, _Immutable, _TypeVarLike, _root=True):
+ """Type variable.
+
+ Usage::
+
+ T = TypeVar('T') # Can be anything
+ A = TypeVar('A', str, bytes) # Must be str or bytes
+
+ Type variables exist primarily for the benefit of static type
+ checkers. They serve as the parameters for generic types as well
+ as for generic function definitions. See class Generic for more
+ information on generic types. Generic functions work as follows:
+
+ def repeat(x: T, n: int) -> List[T]:
+ '''Return a list containing n references to x.'''
+ return [x]*n
+
+ def longest(x: A, y: A) -> A:
+ '''Return the longest of two strings.'''
+ return x if len(x) >= len(y) else y
+
+ The latter example's signature is essentially the overloading
+ of (str, str) -> str and (bytes, bytes) -> bytes. Also note
+ that if the arguments are instances of some subclass of str,
+ the return type is still plain str.
+
+ At runtime, isinstance(x, T) and issubclass(C, T) will raise TypeError.
+
+ Type variables defined with covariant=True or contravariant=True
+ can be used to declare covariant or contravariant generic types.
+ See PEP 484 for more details. By default generic types are invariant
+ in all type variables.
+
+ Type variables can be introspected. e.g.:
+
+ T.__name__ == 'T'
+ T.__constraints__ == ()
+ T.__covariant__ == False
+ T.__contravariant__ = False
+ A.__constraints__ == (str, bytes)
+
+ Note that only type variables defined in global scope can be pickled.
+ """
+
+ __slots__ = ('__name__', '__bound__', '__constraints__',
+ '__covariant__', '__contravariant__', '__dict__')
+
+ def __init__(self, name, *constraints, bound=None,
+ covariant=False, contravariant=False):
+ self.__name__ = name
+ super().__init__(bound, covariant, contravariant)
+ if constraints and bound is not None:
+ raise TypeError("Constraints cannot be combined with bound=...")
+ if constraints and len(constraints) == 1:
+ raise TypeError("A single constraint is not allowed")
+ msg = "TypeVar(name, constraint, ...): constraints must be types."
+ self.__constraints__ = tuple(_type_check(t, msg) for t in constraints)
+ try:
+ def_mod = sys._getframe(1).f_globals.get('__name__', '__main__') # for pickling
+ except (AttributeError, ValueError):
+ def_mod = None
+ if def_mod != 'typing':
+ self.__module__ = def_mod
+
+
+class ParamSpecArgs(_Final, _Immutable, _root=True):
+ """The args for a ParamSpec object.
+
+ Given a ParamSpec object P, P.args is an instance of ParamSpecArgs.
+
+ ParamSpecArgs objects have a reference back to their ParamSpec:
+
+ P.args.__origin__ is P
+
+ This type is meant for runtime introspection and has no special meaning to
+ static type checkers.
+ """
+ def __init__(self, origin):
+ self.__origin__ = origin
+
+ def __repr__(self):
+ return f"{self.__origin__.__name__}.args"
+
+ def __eq__(self, other):
+ if not isinstance(other, ParamSpecArgs):
+ return NotImplemented
+ return self.__origin__ == other.__origin__
+
+
+class ParamSpecKwargs(_Final, _Immutable, _root=True):
+ """The kwargs for a ParamSpec object.
+
+ Given a ParamSpec object P, P.kwargs is an instance of ParamSpecKwargs.
+
+ ParamSpecKwargs objects have a reference back to their ParamSpec:
+
+ P.kwargs.__origin__ is P
+
+ This type is meant for runtime introspection and has no special meaning to
+ static type checkers.
+ """
+ def __init__(self, origin):
+ self.__origin__ = origin
+
+ def __repr__(self):
+ return f"{self.__origin__.__name__}.kwargs"
+
+ def __eq__(self, other):
+ if not isinstance(other, ParamSpecKwargs):
+ return NotImplemented
+ return self.__origin__ == other.__origin__
+
+
+class ParamSpec(_Final, _Immutable, _TypeVarLike, _root=True):
+ """Parameter specification variable.
+
+ Usage::
+
+ P = ParamSpec('P')
+
+ Parameter specification variables exist primarily for the benefit of static
+ type checkers. They are used to forward the parameter types of one
+ callable to another callable, a pattern commonly found in higher order
+ functions and decorators. They are only valid when used in ``Concatenate``,
+ or as the first argument to ``Callable``, or as parameters for user-defined
+ Generics. See class Generic for more information on generic types. An
+ example for annotating a decorator::
+
+ T = TypeVar('T')
+ P = ParamSpec('P')
+
+ def add_logging(f: Callable[P, T]) -> Callable[P, T]:
+ '''A type-safe decorator to add logging to a function.'''
+ def inner(*args: P.args, **kwargs: P.kwargs) -> T:
+ logging.info(f'{f.__name__} was called')
+ return f(*args, **kwargs)
+ return inner
+
+ @add_logging
+ def add_two(x: float, y: float) -> float:
+ '''Add two numbers together.'''
+ return x + y
+
+ Parameter specification variables defined with covariant=True or
+ contravariant=True can be used to declare covariant or contravariant
+ generic types. These keyword arguments are valid, but their actual semantics
+ are yet to be decided. See PEP 612 for details.
+
+ Parameter specification variables can be introspected. e.g.:
+
+ P.__name__ == 'P'
+ P.__bound__ == None
+ P.__covariant__ == False
+ P.__contravariant__ == False
+
+ Note that only parameter specification variables defined in global scope can
+ be pickled.
+ """
+
+ __slots__ = ('__name__', '__bound__', '__covariant__', '__contravariant__',
+ '__dict__')
+
+ @property
+ def args(self):
+ return ParamSpecArgs(self)
+
+ @property
+ def kwargs(self):
+ return ParamSpecKwargs(self)
+
+ def __init__(self, name, *, bound=None, covariant=False, contravariant=False):
+ self.__name__ = name
+ super().__init__(bound, covariant, contravariant)
+ try:
+ def_mod = sys._getframe(1).f_globals.get('__name__', '__main__')
+ except (AttributeError, ValueError):
+ def_mod = None
+ if def_mod != 'typing':
+ self.__module__ = def_mod
+
+
+def _is_dunder(attr):
+ return attr.startswith('__') and attr.endswith('__')
+
+class _BaseGenericAlias(_Final, _root=True):
+ """The central part of internal API.
+
+ This represents a generic version of type 'origin' with type arguments 'params'.
+ There are two kind of these aliases: user defined and special. The special ones
+ are wrappers around builtin collections and ABCs in collections.abc. These must
+ have 'name' always set. If 'inst' is False, then the alias can't be instantiated,
+ this is used by e.g. typing.List and typing.Dict.
+ """
+ def __init__(self, origin, *, inst=True, name=None):
+ self._inst = inst
+ self._name = name
+ self.__origin__ = origin
+ self.__slots__ = None # This is not documented.
+
+ def __call__(self, *args, **kwargs):
+ if not self._inst:
+ raise TypeError(f"Type {self._name} cannot be instantiated; "
+ f"use {self.__origin__.__name__}() instead")
+ result = self.__origin__(*args, **kwargs)
+ try:
+ result.__orig_class__ = self
+ except AttributeError:
+ pass
+ return result
+
+ def __mro_entries__(self, bases):
+ res = []
+ if self.__origin__ not in bases:
+ res.append(self.__origin__)
+ i = bases.index(self)
+ for b in bases[i+1:]:
+ if isinstance(b, _BaseGenericAlias) or issubclass(b, Generic):
+ break
+ else:
+ res.append(Generic)
+ return tuple(res)
+
+ def __getattr__(self, attr):
+ if attr in {'__name__', '__qualname__'}:
+ return self._name or self.__origin__.__name__
+
+ # We are careful for copy and pickle.
+ # Also for simplicity we don't relay any dunder names
+ if '__origin__' in self.__dict__ and not _is_dunder(attr):
+ return getattr(self.__origin__, attr)
+ raise AttributeError(attr)
+
+ def __setattr__(self, attr, val):
+ if _is_dunder(attr) or attr in {'_name', '_inst', '_nparams',
+ '_typevar_types', '_paramspec_tvars'}:
+ super().__setattr__(attr, val)
+ else:
+ setattr(self.__origin__, attr, val)
+
+ def __instancecheck__(self, obj):
+ return self.__subclasscheck__(type(obj))
+
+ def __subclasscheck__(self, cls):
+ raise TypeError("Subscripted generics cannot be used with"
+ " class and instance checks")
+
+ def __dir__(self):
+ return list(set(super().__dir__()
+ + [attr for attr in dir(self.__origin__) if not _is_dunder(attr)]))
+
+# Special typing constructs Union, Optional, Generic, Callable and Tuple
+# use three special attributes for internal bookkeeping of generic types:
+# * __parameters__ is a tuple of unique free type parameters of a generic
+# type, for example, Dict[T, T].__parameters__ == (T,);
+# * __origin__ keeps a reference to a type that was subscripted,
+# e.g., Union[T, int].__origin__ == Union, or the non-generic version of
+# the type.
+# * __args__ is a tuple of all arguments used in subscripting,
+# e.g., Dict[T, int].__args__ == (T, int).
+
+
+class _GenericAlias(_BaseGenericAlias, _root=True):
+ def __init__(self, origin, params, *, inst=True, name=None,
+ _typevar_types=TypeVar,
+ _paramspec_tvars=False):
+ super().__init__(origin, inst=inst, name=name)
+ if not isinstance(params, tuple):
+ params = (params,)
+ self.__args__ = tuple(... if a is _TypingEllipsis else
+ () if a is _TypingEmpty else
+ a for a in params)
+ self.__parameters__ = _collect_type_vars(params, typevar_types=_typevar_types)
+ self._typevar_types = _typevar_types
+ self._paramspec_tvars = _paramspec_tvars
+ if not name:
+ self.__module__ = origin.__module__
+
+ def __eq__(self, other):
+ if not isinstance(other, _GenericAlias):
+ return NotImplemented
+ return (self.__origin__ == other.__origin__
+ and self.__args__ == other.__args__)
+
+ def __hash__(self):
+ return hash((self.__origin__, self.__args__))
+
+ def __or__(self, right):
+ return Union[self, right]
+
+ def __ror__(self, left):
+ return Union[left, self]
+
+ @_tp_cache
+ def __getitem__(self, params):
+ if self.__origin__ in (Generic, Protocol):
+ # Can't subscript Generic[...] or Protocol[...].
+ raise TypeError(f"Cannot subscript already-subscripted {self}")
+ if not isinstance(params, tuple):
+ params = (params,)
+ params = tuple(_type_convert(p) for p in params)
+ if (self._paramspec_tvars
+ and any(isinstance(t, ParamSpec) for t in self.__parameters__)):
+ params = _prepare_paramspec_params(self, params)
+ else:
+ _check_generic(self, params, len(self.__parameters__))
+
+ subst = dict(zip(self.__parameters__, params))
+ new_args = []
+ for arg in self.__args__:
+ if isinstance(arg, self._typevar_types):
+ if isinstance(arg, ParamSpec):
+ arg = subst[arg]
+ if not _is_param_expr(arg):
+ raise TypeError(f"Expected a list of types, an ellipsis, "
+ f"ParamSpec, or Concatenate. Got {arg}")
+ else:
+ arg = subst[arg]
+ elif isinstance(arg, (_GenericAlias, GenericAlias, types.UnionType)):
+ subparams = arg.__parameters__
+ if subparams:
+ subargs = tuple(subst[x] for x in subparams)
+ arg = arg[subargs]
+ # Required to flatten out the args for CallableGenericAlias
+ if self.__origin__ == collections.abc.Callable and isinstance(arg, tuple):
+ new_args.extend(arg)
+ else:
+ new_args.append(arg)
+ return self.copy_with(tuple(new_args))
+
+ def copy_with(self, params):
+ return self.__class__(self.__origin__, params, name=self._name, inst=self._inst,
+ _typevar_types=self._typevar_types,
+ _paramspec_tvars=self._paramspec_tvars)
+
+ def __repr__(self):
+ if self._name:
+ name = 'typing.' + self._name
+ else:
+ name = _type_repr(self.__origin__)
+ args = ", ".join([_type_repr(a) for a in self.__args__])
+ return f'{name}[{args}]'
+
+ def __reduce__(self):
+ if self._name:
+ origin = globals()[self._name]
+ else:
+ origin = self.__origin__
+ args = tuple(self.__args__)
+ if len(args) == 1 and (not isinstance(args[0], tuple) or
+ origin is Tuple and not args[0]):
+ args, = args
+ return operator.getitem, (origin, args)
+
+ def __mro_entries__(self, bases):
+ if isinstance(self.__origin__, _SpecialForm):
+ raise TypeError(f"Cannot subclass {self!r}")
+
+ if self._name: # generic version of an ABC or built-in class
+ return super().__mro_entries__(bases)
+ if self.__origin__ is Generic:
+ if Protocol in bases:
+ return ()
+ i = bases.index(self)
+ for b in bases[i+1:]:
+ if isinstance(b, _BaseGenericAlias) and b is not self:
+ return ()
+ return (self.__origin__,)
+
+
+# _nparams is the number of accepted parameters, e.g. 0 for Hashable,
+# 1 for List and 2 for Dict. It may be -1 if variable number of
+# parameters are accepted (needs custom __getitem__).
+
+class _SpecialGenericAlias(_BaseGenericAlias, _root=True):
+ def __init__(self, origin, nparams, *, inst=True, name=None):
+ if name is None:
+ name = origin.__name__
+ super().__init__(origin, inst=inst, name=name)
+ self._nparams = nparams
+ if origin.__module__ == 'builtins':
+ self.__doc__ = f'A generic version of {origin.__qualname__}.'
+ else:
+ self.__doc__ = f'A generic version of {origin.__module__}.{origin.__qualname__}.'
+
+ @_tp_cache
+ def __getitem__(self, params):
+ if not isinstance(params, tuple):
+ params = (params,)
+ msg = "Parameters to generic types must be types."
+ params = tuple(_type_check(p, msg) for p in params)
+ _check_generic(self, params, self._nparams)
+ return self.copy_with(params)
+
+ def copy_with(self, params):
+ return _GenericAlias(self.__origin__, params,
+ name=self._name, inst=self._inst)
+
+ def __repr__(self):
+ return 'typing.' + self._name
+
+ def __subclasscheck__(self, cls):
+ if isinstance(cls, _SpecialGenericAlias):
+ return issubclass(cls.__origin__, self.__origin__)
+ if not isinstance(cls, _GenericAlias):
+ return issubclass(cls, self.__origin__)
+ return super().__subclasscheck__(cls)
+
+ def __reduce__(self):
+ return self._name
+
+ def __or__(self, right):
+ return Union[self, right]
+
+ def __ror__(self, left):
+ return Union[left, self]
+
+class _CallableGenericAlias(_GenericAlias, _root=True):
+ def __repr__(self):
+ assert self._name == 'Callable'
+ args = self.__args__
+ if len(args) == 2 and _is_param_expr(args[0]):
+ return super().__repr__()
+ return (f'typing.Callable'
+ f'[[{", ".join([_type_repr(a) for a in args[:-1]])}], '
+ f'{_type_repr(args[-1])}]')
+
+ def __reduce__(self):
+ args = self.__args__
+ if not (len(args) == 2 and _is_param_expr(args[0])):
+ args = list(args[:-1]), args[-1]
+ return operator.getitem, (Callable, args)
+
+
+class _CallableType(_SpecialGenericAlias, _root=True):
+ def copy_with(self, params):
+ return _CallableGenericAlias(self.__origin__, params,
+ name=self._name, inst=self._inst,
+ _typevar_types=(TypeVar, ParamSpec),
+ _paramspec_tvars=True)
+
+ def __getitem__(self, params):
+ if not isinstance(params, tuple) or len(params) != 2:
+ raise TypeError("Callable must be used as "
+ "Callable[[arg, ...], result].")
+ args, result = params
+ # This relaxes what args can be on purpose to allow things like
+ # PEP 612 ParamSpec. Responsibility for whether a user is using
+ # Callable[...] properly is deferred to static type checkers.
+ if isinstance(args, list):
+ params = (tuple(args), result)
+ else:
+ params = (args, result)
+ return self.__getitem_inner__(params)
+
+ @_tp_cache
+ def __getitem_inner__(self, params):
+ args, result = params
+ msg = "Callable[args, result]: result must be a type."
+ result = _type_check(result, msg)
+ if args is Ellipsis:
+ return self.copy_with((_TypingEllipsis, result))
+ if not isinstance(args, tuple):
+ args = (args,)
+ args = tuple(_type_convert(arg) for arg in args)
+ params = args + (result,)
+ return self.copy_with(params)
+
+
+class _TupleType(_SpecialGenericAlias, _root=True):
+ @_tp_cache
+ def __getitem__(self, params):
+ if params == ():
+ return self.copy_with((_TypingEmpty,))
+ if not isinstance(params, tuple):
+ params = (params,)
+ if len(params) == 2 and params[1] is ...:
+ msg = "Tuple[t, ...]: t must be a type."
+ p = _type_check(params[0], msg)
+ return self.copy_with((p, _TypingEllipsis))
+ msg = "Tuple[t0, t1, ...]: each t must be a type."
+ params = tuple(_type_check(p, msg) for p in params)
+ return self.copy_with(params)
+
+
+class _UnionGenericAlias(_GenericAlias, _root=True):
+ def copy_with(self, params):
+ return Union[params]
+
+ def __eq__(self, other):
+ if not isinstance(other, (_UnionGenericAlias, types.UnionType)):
+ return NotImplemented
+ return set(self.__args__) == set(other.__args__)
+
+ def __hash__(self):
+ return hash(frozenset(self.__args__))
+
+ def __repr__(self):
+ args = self.__args__
+ if len(args) == 2:
+ if args[0] is type(None):
+ return f'typing.Optional[{_type_repr(args[1])}]'
+ elif args[1] is type(None):
+ return f'typing.Optional[{_type_repr(args[0])}]'
+ return super().__repr__()
+
+ def __instancecheck__(self, obj):
+ return self.__subclasscheck__(type(obj))
+
+ def __subclasscheck__(self, cls):
+ for arg in self.__args__:
+ if issubclass(cls, arg):
+ return True
+
+ def __reduce__(self):
+ func, (origin, args) = super().__reduce__()
+ return func, (Union, args)
+
+
+def _value_and_type_iter(parameters):
+ return ((p, type(p)) for p in parameters)
+
+
+class _LiteralGenericAlias(_GenericAlias, _root=True):
+
+ def __eq__(self, other):
+ if not isinstance(other, _LiteralGenericAlias):
+ return NotImplemented
+
+ return set(_value_and_type_iter(self.__args__)) == set(_value_and_type_iter(other.__args__))
+
+ def __hash__(self):
+ return hash(frozenset(_value_and_type_iter(self.__args__)))
+
+
+class _ConcatenateGenericAlias(_GenericAlias, _root=True):
+ def copy_with(self, params):
+ if isinstance(params[-1], (list, tuple)):
+ return (*params[:-1], *params[-1])
+ if isinstance(params[-1], _ConcatenateGenericAlias):
+ params = (*params[:-1], *params[-1].__args__)
+ elif not isinstance(params[-1], ParamSpec):
+ raise TypeError("The last parameter to Concatenate should be a "
+ "ParamSpec variable.")
+ return super().copy_with(params)
+
+
+class Generic:
+ """Abstract base class for generic types.
+
+ A generic type is typically declared by inheriting from
+ this class parameterized with one or more type variables.
+ For example, a generic mapping type might be defined as::
+
+ class Mapping(Generic[KT, VT]):
+ def __getitem__(self, key: KT) -> VT:
+ ...
+ # Etc.
+
+ This class can then be used as follows::
+
+ def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
+ try:
+ return mapping[key]
+ except KeyError:
+ return default
+ """
+ __slots__ = ()
+ _is_protocol = False
+
+ @_tp_cache
+ def __class_getitem__(cls, params):
+ if not isinstance(params, tuple):
+ params = (params,)
+ if not params and cls is not Tuple:
+ raise TypeError(
+ f"Parameter list to {cls.__qualname__}[...] cannot be empty")
+ params = tuple(_type_convert(p) for p in params)
+ if cls in (Generic, Protocol):
+ # Generic and Protocol can only be subscripted with unique type variables.
+ if not all(isinstance(p, (TypeVar, ParamSpec)) for p in params):
+ raise TypeError(
+ f"Parameters to {cls.__name__}[...] must all be type variables "
+ f"or parameter specification variables.")
+ if len(set(params)) != len(params):
+ raise TypeError(
+ f"Parameters to {cls.__name__}[...] must all be unique")
+ else:
+ # Subscripting a regular Generic subclass.
+ if any(isinstance(t, ParamSpec) for t in cls.__parameters__):
+ params = _prepare_paramspec_params(cls, params)
+ else:
+ _check_generic(cls, params, len(cls.__parameters__))
+ return _GenericAlias(cls, params,
+ _typevar_types=(TypeVar, ParamSpec),
+ _paramspec_tvars=True)
+
+ def __init_subclass__(cls, *args, **kwargs):
+ super().__init_subclass__(*args, **kwargs)
+ tvars = []
+ if '__orig_bases__' in cls.__dict__:
+ error = Generic in cls.__orig_bases__
+ else:
+ error = Generic in cls.__bases__ and cls.__name__ != 'Protocol'
+ if error:
+ raise TypeError("Cannot inherit from plain Generic")
+ if '__orig_bases__' in cls.__dict__:
+ tvars = _collect_type_vars(cls.__orig_bases__, (TypeVar, ParamSpec))
+ # Look for Generic[T1, ..., Tn].
+ # If found, tvars must be a subset of it.
+ # If not found, tvars is it.
+ # Also check for and reject plain Generic,
+ # and reject multiple Generic[...].
+ gvars = None
+ for base in cls.__orig_bases__:
+ if (isinstance(base, _GenericAlias) and
+ base.__origin__ is Generic):
+ if gvars is not None:
+ raise TypeError(
+ "Cannot inherit from Generic[...] multiple types.")
+ gvars = base.__parameters__
+ if gvars is not None:
+ tvarset = set(tvars)
+ gvarset = set(gvars)
+ if not tvarset <= gvarset:
+ s_vars = ', '.join(str(t) for t in tvars if t not in gvarset)
+ s_args = ', '.join(str(g) for g in gvars)
+ raise TypeError(f"Some type variables ({s_vars}) are"
+ f" not listed in Generic[{s_args}]")
+ tvars = gvars
+ cls.__parameters__ = tuple(tvars)
+
+
+class _TypingEmpty:
+ """Internal placeholder for () or []. Used by TupleMeta and CallableMeta
+ to allow empty list/tuple in specific places, without allowing them
+ to sneak in where prohibited.
+ """
+
+
+class _TypingEllipsis:
+ """Internal placeholder for ... (ellipsis)."""
+
+
+_TYPING_INTERNALS = ['__parameters__', '__orig_bases__', '__orig_class__',
+ '_is_protocol', '_is_runtime_protocol']
+
+_SPECIAL_NAMES = ['__abstractmethods__', '__annotations__', '__dict__', '__doc__',
+ '__init__', '__module__', '__new__', '__slots__',
+ '__subclasshook__', '__weakref__', '__class_getitem__']
+
+# These special attributes will be not collected as protocol members.
+EXCLUDED_ATTRIBUTES = _TYPING_INTERNALS + _SPECIAL_NAMES + ['_MutableMapping__marker']
+
+
+def _get_protocol_attrs(cls):
+ """Collect protocol members from a protocol class objects.
+
+ This includes names actually defined in the class dictionary, as well
+ as names that appear in annotations. Special names (above) are skipped.
+ """
+ attrs = set()
+ for base in cls.__mro__[:-1]: # without object
+ if base.__name__ in ('Protocol', 'Generic'):
+ continue
+ annotations = getattr(base, '__annotations__', {})
+ for attr in list(base.__dict__.keys()) + list(annotations.keys()):
+ if not attr.startswith('_abc_') and attr not in EXCLUDED_ATTRIBUTES:
+ attrs.add(attr)
+ return attrs
+
+
+def _is_callable_members_only(cls):
+ # PEP 544 prohibits using issubclass() with protocols that have non-method members.
+ return all(callable(getattr(cls, attr, None)) for attr in _get_protocol_attrs(cls))
+
+
+def _no_init_or_replace_init(self, *args, **kwargs):
+ cls = type(self)
+
+ if cls._is_protocol:
+ raise TypeError('Protocols cannot be instantiated')
+
+ # Already using a custom `__init__`. No need to calculate correct
+ # `__init__` to call. This can lead to RecursionError. See bpo-45121.
+ if cls.__init__ is not _no_init_or_replace_init:
+ return
+
+ # Initially, `__init__` of a protocol subclass is set to `_no_init_or_replace_init`.
+ # The first instantiation of the subclass will call `_no_init_or_replace_init` which
+ # searches for a proper new `__init__` in the MRO. The new `__init__`
+ # replaces the subclass' old `__init__` (ie `_no_init_or_replace_init`). Subsequent
+ # instantiation of the protocol subclass will thus use the new
+ # `__init__` and no longer call `_no_init_or_replace_init`.
+ for base in cls.__mro__:
+ init = base.__dict__.get('__init__', _no_init_or_replace_init)
+ if init is not _no_init_or_replace_init:
+ cls.__init__ = init
+ break
+ else:
+ # should not happen
+ cls.__init__ = object.__init__
+
+ cls.__init__(self, *args, **kwargs)
+
+
+def _caller(depth=1, default='__main__'):
+ try:
+ return sys._getframe(depth + 1).f_globals.get('__name__', default)
+ except (AttributeError, ValueError): # For platforms without _getframe()
+ return None
+
+
+def _allow_reckless_class_checks(depth=3):
+ """Allow instance and class checks for special stdlib modules.
+
+ The abc and functools modules indiscriminately call isinstance() and
+ issubclass() on the whole MRO of a user class, which may contain protocols.
+ """
+ try:
+ return sys._getframe(depth).f_globals['__name__'] in ['abc', 'functools']
+ except (AttributeError, ValueError): # For platforms without _getframe().
+ return True
+
+
+_PROTO_ALLOWLIST = {
+ 'collections.abc': [
+ 'Callable', 'Awaitable', 'Iterable', 'Iterator', 'AsyncIterable',
+ 'Hashable', 'Sized', 'Container', 'Collection', 'Reversible',
+ ],
+ 'contextlib': ['AbstractContextManager', 'AbstractAsyncContextManager'],
+}
+
+
+class _ProtocolMeta(ABCMeta):
+ # This metaclass is really unfortunate and exists only because of
+ # the lack of __instancehook__.
+ def __instancecheck__(cls, instance):
+ # We need this method for situations where attributes are
+ # assigned in __init__.
+ if (
+ getattr(cls, '_is_protocol', False) and
+ not getattr(cls, '_is_runtime_protocol', False) and
+ not _allow_reckless_class_checks(depth=2)
+ ):
+ raise TypeError("Instance and class checks can only be used with"
+ " @runtime_checkable protocols")
+
+ if ((not getattr(cls, '_is_protocol', False) or
+ _is_callable_members_only(cls)) and
+ issubclass(instance.__class__, cls)):
+ return True
+ if cls._is_protocol:
+ if all(hasattr(instance, attr) and
+ # All *methods* can be blocked by setting them to None.
+ (not callable(getattr(cls, attr, None)) or
+ getattr(instance, attr) is not None)
+ for attr in _get_protocol_attrs(cls)):
+ return True
+ return super().__instancecheck__(instance)
+
+
+class Protocol(Generic, metaclass=_ProtocolMeta):
+ """Base class for protocol classes.
+
+ Protocol classes are defined as::
+
+ class Proto(Protocol):
+ def meth(self) -> int:
+ ...
+
+ Such classes are primarily used with static type checkers that recognize
+ structural subtyping (static duck-typing), for example::
+
+ class C:
+ def meth(self) -> int:
+ return 0
+
+ def func(x: Proto) -> int:
+ return x.meth()
+
+ func(C()) # Passes static type check
+
+ See PEP 544 for details. Protocol classes decorated with
+ @typing.runtime_checkable act as simple-minded runtime protocols that check
+ only the presence of given attributes, ignoring their type signatures.
+ Protocol classes can be generic, they are defined as::
+
+ class GenProto(Protocol[T]):
+ def meth(self) -> T:
+ ...
+ """
+ __slots__ = ()
+ _is_protocol = True
+ _is_runtime_protocol = False
+
+ def __init_subclass__(cls, *args, **kwargs):
+ super().__init_subclass__(*args, **kwargs)
+
+ # Determine if this is a protocol or a concrete subclass.
+ if not cls.__dict__.get('_is_protocol', False):
+ cls._is_protocol = any(b is Protocol for b in cls.__bases__)
+
+ # Set (or override) the protocol subclass hook.
+ def _proto_hook(other):
+ if not cls.__dict__.get('_is_protocol', False):
+ return NotImplemented
+
+ # First, perform various sanity checks.
+ if not getattr(cls, '_is_runtime_protocol', False):
+ if _allow_reckless_class_checks():
+ return NotImplemented
+ raise TypeError("Instance and class checks can only be used with"
+ " @runtime_checkable protocols")
+ if not _is_callable_members_only(cls):
+ if _allow_reckless_class_checks():
+ return NotImplemented
+ raise TypeError("Protocols with non-method members"
+ " don't support issubclass()")
+ if not isinstance(other, type):
+ # Same error message as for issubclass(1, int).
+ raise TypeError('issubclass() arg 1 must be a class')
+
+ # Second, perform the actual structural compatibility check.
+ for attr in _get_protocol_attrs(cls):
+ for base in other.__mro__:
+ # Check if the members appears in the class dictionary...
+ if attr in base.__dict__:
+ if base.__dict__[attr] is None:
+ return NotImplemented
+ break
+
+ # ...or in annotations, if it is a sub-protocol.
+ annotations = getattr(base, '__annotations__', {})
+ if (isinstance(annotations, collections.abc.Mapping) and
+ attr in annotations and
+ issubclass(other, Generic) and other._is_protocol):
+ break
+ else:
+ return NotImplemented
+ return True
+
+ if '__subclasshook__' not in cls.__dict__:
+ cls.__subclasshook__ = _proto_hook
+
+ # We have nothing more to do for non-protocols...
+ if not cls._is_protocol:
+ return
+
+ # ... otherwise check consistency of bases, and prohibit instantiation.
+ for base in cls.__bases__:
+ if not (base in (object, Generic) or
+ base.__module__ in _PROTO_ALLOWLIST and
+ base.__name__ in _PROTO_ALLOWLIST[base.__module__] or
+ issubclass(base, Generic) and base._is_protocol):
+ raise TypeError('Protocols can only inherit from other'
+ ' protocols, got %r' % base)
+ cls.__init__ = _no_init_or_replace_init
+
+
+class _AnnotatedAlias(_GenericAlias, _root=True):
+ """Runtime representation of an annotated type.
+
+ At its core 'Annotated[t, dec1, dec2, ...]' is an alias for the type 't'
+ with extra annotations. The alias behaves like a normal typing alias,
+ instantiating is the same as instantiating the underlying type, binding
+ it to types is also the same.
+ """
+ def __init__(self, origin, metadata):
+ if isinstance(origin, _AnnotatedAlias):
+ metadata = origin.__metadata__ + metadata
+ origin = origin.__origin__
+ super().__init__(origin, origin)
+ self.__metadata__ = metadata
+
+ def copy_with(self, params):
+ assert len(params) == 1
+ new_type = params[0]
+ return _AnnotatedAlias(new_type, self.__metadata__)
+
+ def __repr__(self):
+ return "typing.Annotated[{}, {}]".format(
+ _type_repr(self.__origin__),
+ ", ".join(repr(a) for a in self.__metadata__)
+ )
+
+ def __reduce__(self):
+ return operator.getitem, (
+ Annotated, (self.__origin__,) + self.__metadata__
+ )
+
+ def __eq__(self, other):
+ if not isinstance(other, _AnnotatedAlias):
+ return NotImplemented
+ return (self.__origin__ == other.__origin__
+ and self.__metadata__ == other.__metadata__)
+
+ def __hash__(self):
+ return hash((self.__origin__, self.__metadata__))
+
+ def __getattr__(self, attr):
+ if attr in {'__name__', '__qualname__'}:
+ return 'Annotated'
+ return super().__getattr__(attr)
+
+
+class Annotated:
+ """Add context specific metadata to a type.
+
+ Example: Annotated[int, runtime_check.Unsigned] indicates to the
+ hypothetical runtime_check module that this type is an unsigned int.
+ Every other consumer of this type can ignore this metadata and treat
+ this type as int.
+
+ The first argument to Annotated must be a valid type.
+
+ Details:
+
+ - It's an error to call `Annotated` with less than two arguments.
+ - Nested Annotated are flattened::
+
+ Annotated[Annotated[T, Ann1, Ann2], Ann3] == Annotated[T, Ann1, Ann2, Ann3]
+
+ - Instantiating an annotated type is equivalent to instantiating the
+ underlying type::
+
+ Annotated[C, Ann1](5) == C(5)
+
+ - Annotated can be used as a generic type alias::
+
+ Optimized = Annotated[T, runtime.Optimize()]
+ Optimized[int] == Annotated[int, runtime.Optimize()]
+
+ OptimizedList = Annotated[List[T], runtime.Optimize()]
+ OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
+ """
+
+ __slots__ = ()
+
+ def __new__(cls, *args, **kwargs):
+ raise TypeError("Type Annotated cannot be instantiated.")
+
+ @_tp_cache
+ def __class_getitem__(cls, params):
+ if not isinstance(params, tuple) or len(params) < 2:
+ raise TypeError("Annotated[...] should be used "
+ "with at least two arguments (a type and an "
+ "annotation).")
+ msg = "Annotated[t, ...]: t must be a type."
+ origin = _type_check(params[0], msg, allow_special_forms=True)
+ metadata = tuple(params[1:])
+ return _AnnotatedAlias(origin, metadata)
+
+ def __init_subclass__(cls, *args, **kwargs):
+ raise TypeError(
+ "Cannot subclass {}.Annotated".format(cls.__module__)
+ )
+
+
+def runtime_checkable(cls):
+ """Mark a protocol class as a runtime protocol.
+
+ Such protocol can be used with isinstance() and issubclass().
+ Raise TypeError if applied to a non-protocol class.
+ This allows a simple-minded structural check very similar to
+ one trick ponies in collections.abc such as Iterable.
+ For example::
+
+ @runtime_checkable
+ class Closable(Protocol):
+ def close(self): ...
+
+ assert isinstance(open('/some/file'), Closable)
+
+ Warning: this will check only the presence of the required methods,
+ not their type signatures!
+ """
+ if not issubclass(cls, Generic) or not cls._is_protocol:
+ raise TypeError('@runtime_checkable can be only applied to protocol classes,'
+ ' got %r' % cls)
+ cls._is_runtime_protocol = True
+ return cls
+
+
+def cast(typ, val):
+ """Cast a value to a type.
+
+ This returns the value unchanged. To the type checker this
+ signals that the return value has the designated type, but at
+ runtime we intentionally don't check anything (we want this
+ to be as fast as possible).
+ """
+ return val
+
+
+def _get_defaults(func):
+ """Internal helper to extract the default arguments, by name."""
+ try:
+ code = func.__code__
+ except AttributeError:
+ # Some built-in functions don't have __code__, __defaults__, etc.
+ return {}
+ pos_count = code.co_argcount
+ arg_names = code.co_varnames
+ arg_names = arg_names[:pos_count]
+ defaults = func.__defaults__ or ()
+ kwdefaults = func.__kwdefaults__
+ res = dict(kwdefaults) if kwdefaults else {}
+ pos_offset = pos_count - len(defaults)
+ for name, value in zip(arg_names[pos_offset:], defaults):
+ assert name not in res
+ res[name] = value
+ return res
+
+
+_allowed_types = (types.FunctionType, types.BuiltinFunctionType,
+ types.MethodType, types.ModuleType,
+ WrapperDescriptorType, MethodWrapperType, MethodDescriptorType)
+
+
+def get_type_hints(obj, globalns=None, localns=None, include_extras=False):
+ """Return type hints for an object.
+
+ This is often the same as obj.__annotations__, but it handles
+ forward references encoded as string literals, adds Optional[t] if a
+ default value equal to None is set and recursively replaces all
+ 'Annotated[T, ...]' with 'T' (unless 'include_extras=True').
+
+ The argument may be a module, class, method, or function. The annotations
+ are returned as a dictionary. For classes, annotations include also
+ inherited members.
+
+ TypeError is raised if the argument is not of a type that can contain
+ annotations, and an empty dictionary is returned if no annotations are
+ present.
+
+ BEWARE -- the behavior of globalns and localns is counterintuitive
+ (unless you are familiar with how eval() and exec() work). The
+ search order is locals first, then globals.
+
+ - If no dict arguments are passed, an attempt is made to use the
+ globals from obj (or the respective module's globals for classes),
+ and these are also used as the locals. If the object does not appear
+ to have globals, an empty dictionary is used. For classes, the search
+ order is globals first then locals.
+
+ - If one dict argument is passed, it is used for both globals and
+ locals.
+
+ - If two dict arguments are passed, they specify globals and
+ locals, respectively.
+ """
+
+ if getattr(obj, '__no_type_check__', None):
+ return {}
+ # Classes require a special treatment.
+ if isinstance(obj, type):
+ hints = {}
+ for base in reversed(obj.__mro__):
+ if globalns is None:
+ base_globals = getattr(sys.modules.get(base.__module__, None), '__dict__', {})
+ else:
+ base_globals = globalns
+ ann = base.__dict__.get('__annotations__', {})
+ if isinstance(ann, types.GetSetDescriptorType):
+ ann = {}
+ base_locals = dict(vars(base)) if localns is None else localns
+ if localns is None and globalns is None:
+ # This is surprising, but required. Before Python 3.10,
+ # get_type_hints only evaluated the globalns of
+ # a class. To maintain backwards compatibility, we reverse
+ # the globalns and localns order so that eval() looks into
+ # *base_globals* first rather than *base_locals*.
+ # This only affects ForwardRefs.
+ base_globals, base_locals = base_locals, base_globals
+ for name, value in ann.items():
+ if value is None:
+ value = type(None)
+ if isinstance(value, str):
+ value = ForwardRef(value, is_argument=False, is_class=True)
+ value = _eval_type(value, base_globals, base_locals)
+ hints[name] = value
+ return hints if include_extras else {k: _strip_annotations(t) for k, t in hints.items()}
+
+ if globalns is None:
+ if isinstance(obj, types.ModuleType):
+ globalns = obj.__dict__
+ else:
+ nsobj = obj
+ # Find globalns for the unwrapped object.
+ while hasattr(nsobj, '__wrapped__'):
+ nsobj = nsobj.__wrapped__
+ globalns = getattr(nsobj, '__globals__', {})
+ if localns is None:
+ localns = globalns
+ elif localns is None:
+ localns = globalns
+ hints = getattr(obj, '__annotations__', None)
+ if hints is None:
+ # Return empty annotations for something that _could_ have them.
+ if isinstance(obj, _allowed_types):
+ return {}
+ else:
+ raise TypeError('{!r} is not a module, class, method, '
+ 'or function.'.format(obj))
+ defaults = _get_defaults(obj)
+ hints = dict(hints)
+ for name, value in hints.items():
+ if value is None:
+ value = type(None)
+ if isinstance(value, str):
+ # class-level forward refs were handled above, this must be either
+ # a module-level annotation or a function argument annotation
+ value = ForwardRef(
+ value,
+ is_argument=not isinstance(obj, types.ModuleType),
+ is_class=False,
+ )
+ value = _eval_type(value, globalns, localns)
+ if name in defaults and defaults[name] is None:
+ value = Optional[value]
+ hints[name] = value
+ return hints if include_extras else {k: _strip_annotations(t) for k, t in hints.items()}
+
+
+def _strip_annotations(t):
+ """Strips the annotations from a given type.
+ """
+ if isinstance(t, _AnnotatedAlias):
+ return _strip_annotations(t.__origin__)
+ if isinstance(t, _GenericAlias):
+ stripped_args = tuple(_strip_annotations(a) for a in t.__args__)
+ if stripped_args == t.__args__:
+ return t
+ return t.copy_with(stripped_args)
+ if isinstance(t, GenericAlias):
+ stripped_args = tuple(_strip_annotations(a) for a in t.__args__)
+ if stripped_args == t.__args__:
+ return t
+ return GenericAlias(t.__origin__, stripped_args)
+ if isinstance(t, types.UnionType):
+ stripped_args = tuple(_strip_annotations(a) for a in t.__args__)
+ if stripped_args == t.__args__:
+ return t
+ return functools.reduce(operator.or_, stripped_args)
+
+ return t
+
+
+def get_origin(tp):
+ """Get the unsubscripted version of a type.
+
+ This supports generic types, Callable, Tuple, Union, Literal, Final, ClassVar
+ and Annotated. Return None for unsupported types. Examples::
+
+ get_origin(Literal[42]) is Literal
+ get_origin(int) is None
+ get_origin(ClassVar[int]) is ClassVar
+ get_origin(Generic) is Generic
+ get_origin(Generic[T]) is Generic
+ get_origin(Union[T, int]) is Union
+ get_origin(List[Tuple[T, T]][int]) == list
+ get_origin(P.args) is P
+ """
+ if isinstance(tp, _AnnotatedAlias):
+ return Annotated
+ if isinstance(tp, (_BaseGenericAlias, GenericAlias,
+ ParamSpecArgs, ParamSpecKwargs)):
+ return tp.__origin__
+ if tp is Generic:
+ return Generic
+ if isinstance(tp, types.UnionType):
+ return types.UnionType
+ return None
+
+
+def get_args(tp):
+ """Get type arguments with all substitutions performed.
+
+ For unions, basic simplifications used by Union constructor are performed.
+ Examples::
+ get_args(Dict[str, int]) == (str, int)
+ get_args(int) == ()
+ get_args(Union[int, Union[T, int], str][int]) == (int, str)
+ get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int])
+ get_args(Callable[[], T][int]) == ([], int)
+ """
+ if isinstance(tp, _AnnotatedAlias):
+ return (tp.__origin__,) + tp.__metadata__
+ if isinstance(tp, (_GenericAlias, GenericAlias)):
+ res = tp.__args__
+ if (tp.__origin__ is collections.abc.Callable
+ and not (len(res) == 2 and _is_param_expr(res[0]))):
+ res = (list(res[:-1]), res[-1])
+ return res
+ if isinstance(tp, types.UnionType):
+ return tp.__args__
+ return ()
+
+
+def is_typeddict(tp):
+ """Check if an annotation is a TypedDict class
+
+ For example::
+ class Film(TypedDict):
+ title: str
+ year: int
+
+ is_typeddict(Film) # => True
+ is_typeddict(Union[list, str]) # => False
+ """
+ return isinstance(tp, _TypedDictMeta)
+
+
+def no_type_check(arg):
+ """Decorator to indicate that annotations are not type hints.
+
+ The argument must be a class or function; if it is a class, it
+ applies recursively to all methods and classes defined in that class
+ (but not to methods defined in its superclasses or subclasses).
+
+ This mutates the function(s) or class(es) in place.
+ """
+ if isinstance(arg, type):
+ arg_attrs = arg.__dict__.copy()
+ for attr, val in arg.__dict__.items():
+ if val in arg.__bases__ + (arg,):
+ arg_attrs.pop(attr)
+ for obj in arg_attrs.values():
+ if isinstance(obj, types.FunctionType):
+ obj.__no_type_check__ = True
+ if isinstance(obj, type):
+ no_type_check(obj)
+ try:
+ arg.__no_type_check__ = True
+ except TypeError: # built-in classes
+ pass
+ return arg
+
+
+def no_type_check_decorator(decorator):
+ """Decorator to give another decorator the @no_type_check effect.
+
+ This wraps the decorator with something that wraps the decorated
+ function in @no_type_check.
+ """
+
+ @functools.wraps(decorator)
+ def wrapped_decorator(*args, **kwds):
+ func = decorator(*args, **kwds)
+ func = no_type_check(func)
+ return func
+
+ return wrapped_decorator
+
+
+def _overload_dummy(*args, **kwds):
+ """Helper for @overload to raise when called."""
+ raise NotImplementedError(
+ "You should not call an overloaded function. "
+ "A series of @overload-decorated functions "
+ "outside a stub module should always be followed "
+ "by an implementation that is not @overload-ed.")
+
+
+def overload(func):
+ """Decorator for overloaded functions/methods.
+
+ In a stub file, place two or more stub definitions for the same
+ function in a row, each decorated with @overload. For example:
+
+ @overload
+ def utf8(value: None) -> None: ...
+ @overload
+ def utf8(value: bytes) -> bytes: ...
+ @overload
+ def utf8(value: str) -> bytes: ...
+
+ In a non-stub file (i.e. a regular .py file), do the same but
+ follow it with an implementation. The implementation should *not*
+ be decorated with @overload. For example:
+
+ @overload
+ def utf8(value: None) -> None: ...
+ @overload
+ def utf8(value: bytes) -> bytes: ...
+ @overload
+ def utf8(value: str) -> bytes: ...
+ def utf8(value):
+ # implementation goes here
+ """
+ return _overload_dummy
+
+
+def final(f):
+ """A decorator to indicate final methods and final classes.
+
+ Use this decorator to indicate to type checkers that the decorated
+ method cannot be overridden, and decorated class cannot be subclassed.
+ For example:
+
+ class Base:
+ @final
+ def done(self) -> None:
+ ...
+ class Sub(Base):
+ def done(self) -> None: # Error reported by type checker
+ ...
+
+ @final
+ class Leaf:
+ ...
+ class Other(Leaf): # Error reported by type checker
+ ...
+
+ There is no runtime checking of these properties.
+ """
+ return f
+
+
+# Some unconstrained type variables. These are used by the container types.
+# (These are not for export.)
+T = TypeVar('T') # Any type.
+KT = TypeVar('KT') # Key type.
+VT = TypeVar('VT') # Value type.
+T_co = TypeVar('T_co', covariant=True) # Any type covariant containers.
+V_co = TypeVar('V_co', covariant=True) # Any type covariant containers.
+VT_co = TypeVar('VT_co', covariant=True) # Value type covariant containers.
+T_contra = TypeVar('T_contra', contravariant=True) # Ditto contravariant.
+# Internal type variable used for Type[].
+CT_co = TypeVar('CT_co', covariant=True, bound=type)
+
+# A useful type variable with constraints. This represents string types.
+# (This one *is* for export!)
+AnyStr = TypeVar('AnyStr', bytes, str)
+
+
+# Various ABCs mimicking those in collections.abc.
+_alias = _SpecialGenericAlias
+
+Hashable = _alias(collections.abc.Hashable, 0) # Not generic.
+Awaitable = _alias(collections.abc.Awaitable, 1)
+Coroutine = _alias(collections.abc.Coroutine, 3)
+AsyncIterable = _alias(collections.abc.AsyncIterable, 1)
+AsyncIterator = _alias(collections.abc.AsyncIterator, 1)
+Iterable = _alias(collections.abc.Iterable, 1)
+Iterator = _alias(collections.abc.Iterator, 1)
+Reversible = _alias(collections.abc.Reversible, 1)
+Sized = _alias(collections.abc.Sized, 0) # Not generic.
+Container = _alias(collections.abc.Container, 1)
+Collection = _alias(collections.abc.Collection, 1)
+Callable = _CallableType(collections.abc.Callable, 2)
+Callable.__doc__ = \
+ """Callable type; Callable[[int], str] is a function of (int) -> str.
+
+ The subscription syntax must always be used with exactly two
+ values: the argument list and the return type. The argument list
+ must be a list of types or ellipsis; the return type must be a single type.
+
+ There is no syntax to indicate optional or keyword arguments,
+ such function types are rarely used as callback types.
+ """
+AbstractSet = _alias(collections.abc.Set, 1, name='AbstractSet')
+MutableSet = _alias(collections.abc.MutableSet, 1)
+# NOTE: Mapping is only covariant in the value type.
+Mapping = _alias(collections.abc.Mapping, 2)
+MutableMapping = _alias(collections.abc.MutableMapping, 2)
+Sequence = _alias(collections.abc.Sequence, 1)
+MutableSequence = _alias(collections.abc.MutableSequence, 1)
+ByteString = _alias(collections.abc.ByteString, 0) # Not generic
+# Tuple accepts variable number of parameters.
+Tuple = _TupleType(tuple, -1, inst=False, name='Tuple')
+Tuple.__doc__ = \
+ """Tuple type; Tuple[X, Y] is the cross-product type of X and Y.
+
+ Example: Tuple[T1, T2] is a tuple of two elements corresponding
+ to type variables T1 and T2. Tuple[int, float, str] is a tuple
+ of an int, a float and a string.
+
+ To specify a variable-length tuple of homogeneous type, use Tuple[T, ...].
+ """
+List = _alias(list, 1, inst=False, name='List')
+Deque = _alias(collections.deque, 1, name='Deque')
+Set = _alias(set, 1, inst=False, name='Set')
+FrozenSet = _alias(frozenset, 1, inst=False, name='FrozenSet')
+MappingView = _alias(collections.abc.MappingView, 1)
+KeysView = _alias(collections.abc.KeysView, 1)
+ItemsView = _alias(collections.abc.ItemsView, 2)
+ValuesView = _alias(collections.abc.ValuesView, 1)
+ContextManager = _alias(contextlib.AbstractContextManager, 1, name='ContextManager')
+AsyncContextManager = _alias(contextlib.AbstractAsyncContextManager, 1, name='AsyncContextManager')
+Dict = _alias(dict, 2, inst=False, name='Dict')
+DefaultDict = _alias(collections.defaultdict, 2, name='DefaultDict')
+OrderedDict = _alias(collections.OrderedDict, 2)
+Counter = _alias(collections.Counter, 1)
+ChainMap = _alias(collections.ChainMap, 2)
+Generator = _alias(collections.abc.Generator, 3)
+AsyncGenerator = _alias(collections.abc.AsyncGenerator, 2)
+Type = _alias(type, 1, inst=False, name='Type')
+Type.__doc__ = \
+ """A special construct usable to annotate class objects.
+
+ For example, suppose we have the following classes::
+
+ class User: ... # Abstract base for User classes
+ class BasicUser(User): ...
+ class ProUser(User): ...
+ class TeamUser(User): ...
+
+ And a function that takes a class argument that's a subclass of
+ User and returns an instance of the corresponding class::
+
+ U = TypeVar('U', bound=User)
+ def new_user(user_class: Type[U]) -> U:
+ user = user_class()
+ # (Here we could write the user object to a database)
+ return user
+
+ joe = new_user(BasicUser)
+
+ At this point the type checker knows that joe has type BasicUser.
+ """
+
+
+@runtime_checkable
+class SupportsInt(Protocol):
+ """An ABC with one abstract method __int__."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __int__(self) -> int:
+ pass
+
+
+@runtime_checkable
+class SupportsFloat(Protocol):
+ """An ABC with one abstract method __float__."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __float__(self) -> float:
+ pass
+
+
+@runtime_checkable
+class SupportsComplex(Protocol):
+ """An ABC with one abstract method __complex__."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __complex__(self) -> complex:
+ pass
+
+
+@runtime_checkable
+class SupportsBytes(Protocol):
+ """An ABC with one abstract method __bytes__."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __bytes__(self) -> bytes:
+ pass
+
+
+@runtime_checkable
+class SupportsIndex(Protocol):
+ """An ABC with one abstract method __index__."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __index__(self) -> int:
+ pass
+
+
+@runtime_checkable
+class SupportsAbs(Protocol[T_co]):
+ """An ABC with one abstract method __abs__ that is covariant in its return type."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __abs__(self) -> T_co:
+ pass
+
+
+@runtime_checkable
+class SupportsRound(Protocol[T_co]):
+ """An ABC with one abstract method __round__ that is covariant in its return type."""
+ __slots__ = ()
+
+ @abstractmethod
+ def __round__(self, ndigits: int = 0) -> T_co:
+ pass
+
+
+def _make_nmtuple(name, types, module, defaults = ()):
+ fields = [n for n, t in types]
+ types = {n: _type_check(t, f"field {n} annotation must be a type")
+ for n, t in types}
+ nm_tpl = collections.namedtuple(name, fields,
+ defaults=defaults, module=module)
+ nm_tpl.__annotations__ = nm_tpl.__new__.__annotations__ = types
+ return nm_tpl
+
+
+# attributes prohibited to set in NamedTuple class syntax
+_prohibited = frozenset({'__new__', '__init__', '__slots__', '__getnewargs__',
+ '_fields', '_field_defaults',
+ '_make', '_replace', '_asdict', '_source'})
+
+_special = frozenset({'__module__', '__name__', '__annotations__'})
+
+
+class NamedTupleMeta(type):
+
+ def __new__(cls, typename, bases, ns):
+ assert bases[0] is _NamedTuple
+ types = ns.get('__annotations__', {})
+ default_names = []
+ for field_name in types:
+ if field_name in ns:
+ default_names.append(field_name)
+ elif default_names:
+ raise TypeError(f"Non-default namedtuple field {field_name} "
+ f"cannot follow default field"
+ f"{'s' if len(default_names) > 1 else ''} "
+ f"{', '.join(default_names)}")
+ nm_tpl = _make_nmtuple(typename, types.items(),
+ defaults=[ns[n] for n in default_names],
+ module=ns['__module__'])
+ # update from user namespace without overriding special namedtuple attributes
+ for key in ns:
+ if key in _prohibited:
+ raise AttributeError("Cannot overwrite NamedTuple attribute " + key)
+ elif key not in _special and key not in nm_tpl._fields:
+ setattr(nm_tpl, key, ns[key])
+ return nm_tpl
+
+
+def NamedTuple(typename, fields=None, /, **kwargs):
+ """Typed version of namedtuple.
+
+ Usage in Python versions >= 3.6::
+
+ class Employee(NamedTuple):
+ name: str
+ id: int
+
+ This is equivalent to::
+
+ Employee = collections.namedtuple('Employee', ['name', 'id'])
+
+ The resulting class has an extra __annotations__ attribute, giving a
+ dict that maps field names to types. (The field names are also in
+ the _fields attribute, which is part of the namedtuple API.)
+ Alternative equivalent keyword syntax is also accepted::
+
+ Employee = NamedTuple('Employee', name=str, id=int)
+
+ In Python versions <= 3.5 use::
+
+ Employee = NamedTuple('Employee', [('name', str), ('id', int)])
+ """
+ if fields is None:
+ fields = kwargs.items()
+ elif kwargs:
+ raise TypeError("Either list of fields or keywords"
+ " can be provided to NamedTuple, not both")
+ try:
+ module = sys._getframe(1).f_globals.get('__name__', '__main__')
+ except (AttributeError, ValueError):
+ module = None
+ return _make_nmtuple(typename, fields, module=module)
+
+_NamedTuple = type.__new__(NamedTupleMeta, 'NamedTuple', (), {})
+
+def _namedtuple_mro_entries(bases):
+ if len(bases) > 1:
+ raise TypeError("Multiple inheritance with NamedTuple is not supported")
+ assert bases[0] is NamedTuple
+ return (_NamedTuple,)
+
+NamedTuple.__mro_entries__ = _namedtuple_mro_entries
+
+
+class _TypedDictMeta(type):
+ def __new__(cls, name, bases, ns, total=True):
+ """Create new typed dict class object.
+
+ This method is called when TypedDict is subclassed,
+ or when TypedDict is instantiated. This way
+ TypedDict supports all three syntax forms described in its docstring.
+ Subclasses and instances of TypedDict return actual dictionaries.
+ """
+ for base in bases:
+ if type(base) is not _TypedDictMeta:
+ raise TypeError('cannot inherit from both a TypedDict type '
+ 'and a non-TypedDict base class')
+ tp_dict = type.__new__(_TypedDictMeta, name, (dict,), ns)
+
+ annotations = {}
+ own_annotations = ns.get('__annotations__', {})
+ own_annotation_keys = set(own_annotations.keys())
+ msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type"
+ own_annotations = {
+ n: _type_check(tp, msg, module=tp_dict.__module__)
+ for n, tp in own_annotations.items()
+ }
+ required_keys = set()
+ optional_keys = set()
+
+ for base in bases:
+ annotations.update(base.__dict__.get('__annotations__', {}))
+ required_keys.update(base.__dict__.get('__required_keys__', ()))
+ optional_keys.update(base.__dict__.get('__optional_keys__', ()))
+
+ annotations.update(own_annotations)
+ if total:
+ required_keys.update(own_annotation_keys)
+ else:
+ optional_keys.update(own_annotation_keys)
+
+ tp_dict.__annotations__ = annotations
+ tp_dict.__required_keys__ = frozenset(required_keys)
+ tp_dict.__optional_keys__ = frozenset(optional_keys)
+ if not hasattr(tp_dict, '__total__'):
+ tp_dict.__total__ = total
+ return tp_dict
+
+ __call__ = dict # static method
+
+ def __subclasscheck__(cls, other):
+ # Typed dicts are only for static structural subtyping.
+ raise TypeError('TypedDict does not support instance and class checks')
+
+ __instancecheck__ = __subclasscheck__
+
+
+def TypedDict(typename, fields=None, /, *, total=True, **kwargs):
+ """A simple typed namespace. At runtime it is equivalent to a plain dict.
+
+ TypedDict creates a dictionary type that expects all of its
+ instances to have a certain set of keys, where each key is
+ associated with a value of a consistent type. This expectation
+ is not checked at runtime but is only enforced by type checkers.
+ Usage::
+
+ class Point2D(TypedDict):
+ x: int
+ y: int
+ label: str
+
+ a: Point2D = {'x': 1, 'y': 2, 'label': 'good'} # OK
+ b: Point2D = {'z': 3, 'label': 'bad'} # Fails type check
+
+ assert Point2D(x=1, y=2, label='first') == dict(x=1, y=2, label='first')
+
+ The type info can be accessed via the Point2D.__annotations__ dict, and
+ the Point2D.__required_keys__ and Point2D.__optional_keys__ frozensets.
+ TypedDict supports two additional equivalent forms::
+
+ Point2D = TypedDict('Point2D', x=int, y=int, label=str)
+ Point2D = TypedDict('Point2D', {'x': int, 'y': int, 'label': str})
+
+ By default, all keys must be present in a TypedDict. It is possible
+ to override this by specifying totality.
+ Usage::
+
+ class point2D(TypedDict, total=False):
+ x: int
+ y: int
+
+ This means that a point2D TypedDict can have any of the keys omitted.A type
+ checker is only expected to support a literal False or True as the value of
+ the total argument. True is the default, and makes all items defined in the
+ class body be required.
+
+ The class syntax is only supported in Python 3.6+, while two other
+ syntax forms work for Python 2.7 and 3.2+
+ """
+ if fields is None:
+ fields = kwargs
+ elif kwargs:
+ raise TypeError("TypedDict takes either a dict or keyword arguments,"
+ " but not both")
+
+ ns = {'__annotations__': dict(fields)}
+ try:
+ # Setting correct module is necessary to make typed dict classes pickleable.
+ ns['__module__'] = sys._getframe(1).f_globals.get('__name__', '__main__')
+ except (AttributeError, ValueError):
+ pass
+
+ return _TypedDictMeta(typename, (), ns, total=total)
+
+_TypedDict = type.__new__(_TypedDictMeta, 'TypedDict', (), {})
+TypedDict.__mro_entries__ = lambda bases: (_TypedDict,)
+
+
+class NewType:
+ """NewType creates simple unique types with almost zero
+ runtime overhead. NewType(name, tp) is considered a subtype of tp
+ by static type checkers. At runtime, NewType(name, tp) returns
+ a dummy callable that simply returns its argument. Usage::
+
+ UserId = NewType('UserId', int)
+
+ def name_by_id(user_id: UserId) -> str:
+ ...
+
+ UserId('user') # Fails type check
+
+ name_by_id(42) # Fails type check
+ name_by_id(UserId(42)) # OK
+
+ num = UserId(5) + 1 # type: int
+ """
+
+ def __init__(self, name, tp):
+ self.__qualname__ = name
+ if '.' in name:
+ name = name.rpartition('.')[-1]
+ self.__name__ = name
+ self.__supertype__ = tp
+ def_mod = _caller()
+ if def_mod != 'typing':
+ self.__module__ = def_mod
+
+ def __repr__(self):
+ return f'{self.__module__}.{self.__qualname__}'
+
+ def __call__(self, x):
+ return x
+
+ def __reduce__(self):
+ return self.__qualname__
+
+ def __or__(self, other):
+ return Union[self, other]
+
+ def __ror__(self, other):
+ return Union[other, self]
+
+
+# Python-version-specific alias (Python 2: unicode; Python 3: str)
+Text = str
+
+
+# Constant that's True when type checking, but False here.
+TYPE_CHECKING = False
+
+
+class IO(Generic[AnyStr]):
+ """Generic base class for TextIO and BinaryIO.
+
+ This is an abstract, generic version of the return of open().
+
+ NOTE: This does not distinguish between the different possible
+ classes (text vs. binary, read vs. write vs. read/write,
+ append-only, unbuffered). The TextIO and BinaryIO subclasses
+ below capture the distinctions between text vs. binary, which is
+ pervasive in the interface; however we currently do not offer a
+ way to track the other distinctions in the type system.
+ """
+
+ __slots__ = ()
+
+ @property
+ @abstractmethod
+ def mode(self) -> str:
+ pass
+
+ @property
+ @abstractmethod
+ def name(self) -> str:
+ pass
+
+ @abstractmethod
+ def close(self) -> None:
+ pass
+
+ @property
+ @abstractmethod
+ def closed(self) -> bool:
+ pass
+
+ @abstractmethod
+ def fileno(self) -> int:
+ pass
+
+ @abstractmethod
+ def flush(self) -> None:
+ pass
+
+ @abstractmethod
+ def isatty(self) -> bool:
+ pass
+
+ @abstractmethod
+ def read(self, n: int = -1) -> AnyStr:
+ pass
+
+ @abstractmethod
+ def readable(self) -> bool:
+ pass
+
+ @abstractmethod
+ def readline(self, limit: int = -1) -> AnyStr:
+ pass
+
+ @abstractmethod
+ def readlines(self, hint: int = -1) -> List[AnyStr]:
+ pass
+
+ @abstractmethod
+ def seek(self, offset: int, whence: int = 0) -> int:
+ pass
+
+ @abstractmethod
+ def seekable(self) -> bool:
+ pass
+
+ @abstractmethod
+ def tell(self) -> int:
+ pass
+
+ @abstractmethod
+ def truncate(self, size: int = None) -> int:
+ pass
+
+ @abstractmethod
+ def writable(self) -> bool:
+ pass
+
+ @abstractmethod
+ def write(self, s: AnyStr) -> int:
+ pass
+
+ @abstractmethod
+ def writelines(self, lines: List[AnyStr]) -> None:
+ pass
+
+ @abstractmethod
+ def __enter__(self) -> 'IO[AnyStr]':
+ pass
+
+ @abstractmethod
+ def __exit__(self, type, value, traceback) -> None:
+ pass
+
+
+class BinaryIO(IO[bytes]):
+ """Typed version of the return of open() in binary mode."""
+
+ __slots__ = ()
+
+ @abstractmethod
+ def write(self, s: Union[bytes, bytearray]) -> int:
+ pass
+
+ @abstractmethod
+ def __enter__(self) -> 'BinaryIO':
+ pass
+
+
+class TextIO(IO[str]):
+ """Typed version of the return of open() in text mode."""
+
+ __slots__ = ()
+
+ @property
+ @abstractmethod
+ def buffer(self) -> BinaryIO:
+ pass
+
+ @property
+ @abstractmethod
+ def encoding(self) -> str:
+ pass
+
+ @property
+ @abstractmethod
+ def errors(self) -> Optional[str]:
+ pass
+
+ @property
+ @abstractmethod
+ def line_buffering(self) -> bool:
+ pass
+
+ @property
+ @abstractmethod
+ def newlines(self) -> Any:
+ pass
+
+ @abstractmethod
+ def __enter__(self) -> 'TextIO':
+ pass
+
+
+class io:
+ """Wrapper namespace for IO generic classes."""
+
+ __all__ = ['IO', 'TextIO', 'BinaryIO']
+ IO = IO
+ TextIO = TextIO
+ BinaryIO = BinaryIO
+
+
+io.__name__ = __name__ + '.io'
+sys.modules[io.__name__] = io
+
+Pattern = _alias(stdlib_re.Pattern, 1)
+Match = _alias(stdlib_re.Match, 1)
+
+class re:
+ """Wrapper namespace for re type aliases."""
+
+ __all__ = ['Pattern', 'Match']
+ Pattern = Pattern
+ Match = Match
+
+
+re.__name__ = __name__ + '.re'
+sys.modules[re.__name__] = re
+Short
+ */ + .o-tooltip--left { + position: relative; + } + + .o-tooltip--left:after { + opacity: 0; + visibility: hidden; + position: absolute; + content: attr(data-tooltip); + padding: .2em; + font-size: .8em; + left: -.2em; + background: grey; + color: white; + white-space: nowrap; + z-index: 2; + border-radius: 2px; + transform: translateX(-102%) translateY(0); + transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); +} + +.o-tooltip--left:hover:after { + display: block; + opacity: 1; + visibility: visible; + transform: translateX(-100%) translateY(0); + transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); + transition-delay: .5s; +} + +/* By default the copy button shouldn't show up when printing a page */ +@media print { + button.copybtn { + display: none; + } +} diff --git a/_static/copybutton.js b/_static/copybutton.js new file mode 100644 index 00000000..2ea7ff3e --- /dev/null +++ b/_static/copybutton.js @@ -0,0 +1,248 @@ +// Localization support +const messages = { + 'en': { + 'copy': 'Copy', + 'copy_to_clipboard': 'Copy to clipboard', + 'copy_success': 'Copied!', + 'copy_failure': 'Failed to copy', + }, + 'es' : { + 'copy': 'Copiar', + 'copy_to_clipboard': 'Copiar al portapapeles', + 'copy_success': '¡Copiado!', + 'copy_failure': 'Error al copiar', + }, + 'de' : { + 'copy': 'Kopieren', + 'copy_to_clipboard': 'In die Zwischenablage kopieren', + 'copy_success': 'Kopiert!', + 'copy_failure': 'Fehler beim Kopieren', + }, + 'fr' : { + 'copy': 'Copier', + 'copy_to_clipboard': 'Copier dans le presse-papier', + 'copy_success': 'Copié !', + 'copy_failure': 'Échec de la copie', + }, + 'ru': { + 'copy': 'Скопировать', + 'copy_to_clipboard': 'Скопировать в буфер', + 'copy_success': 'Скопировано!', + 'copy_failure': 'Не удалось скопировать', + }, + 'zh-CN': { + 'copy': '复制', + 'copy_to_clipboard': '复制到剪贴板', + 'copy_success': '复制成功!', + 'copy_failure': '复制失败', + }, + 'it' : { + 'copy': 'Copiare', + 'copy_to_clipboard': 'Copiato negli appunti', + 'copy_success': 'Copiato!', + 'copy_failure': 'Errore durante la copia', + } +} + +let locale = 'en' +if( document.documentElement.lang !== undefined + && messages[document.documentElement.lang] !== undefined ) { + locale = document.documentElement.lang +} + +let doc_url_root = DOCUMENTATION_OPTIONS.URL_ROOT; +if (doc_url_root == '#') { + doc_url_root = ''; +} + +/** + * SVG files for our copy buttons + */ +let iconCheck = `` + +// If the user specified their own SVG use that, otherwise use the default +let iconCopy = ``; +if (!iconCopy) { + iconCopy = `` +} + +/** + * Set up copy/paste for code blocks + */ + +const runWhenDOMLoaded = cb => { + if (document.readyState != 'loading') { + cb() + } else if (document.addEventListener) { + document.addEventListener('DOMContentLoaded', cb) + } else { + document.attachEvent('onreadystatechange', function() { + if (document.readyState == 'complete') cb() + }) + } +} + +const codeCellId = index => `codecell${index}` + +// Clears selected text since ClipboardJS will select the text when copying +const clearSelection = () => { + if (window.getSelection) { + window.getSelection().removeAllRanges() + } else if (document.selection) { + document.selection.empty() + } +} + +// Changes tooltip text for a moment, then changes it back +// We want the timeout of our `success` class to be a bit shorter than the +// tooltip and icon change, so that we can hide the icon before changing back. +var timeoutIcon = 2000; +var timeoutSuccessClass = 1500; + +const temporarilyChangeTooltip = (el, oldText, newText) => { + el.setAttribute('data-tooltip', newText) + el.classList.add('success') + // Remove success a little bit sooner than we change the tooltip + // So that we can use CSS to hide the copybutton first + setTimeout(() => el.classList.remove('success'), timeoutSuccessClass) + setTimeout(() => el.setAttribute('data-tooltip', oldText), timeoutIcon) +} + +// Changes the copy button icon for two seconds, then changes it back +const temporarilyChangeIcon = (el) => { + el.innerHTML = iconCheck; + setTimeout(() => {el.innerHTML = iconCopy}, timeoutIcon) +} + +const addCopyButtonToCodeCells = () => { + // If ClipboardJS hasn't loaded, wait a bit and try again. This + // happens because we load ClipboardJS asynchronously. + if (window.ClipboardJS === undefined) { + setTimeout(addCopyButtonToCodeCells, 250) + return + } + + // Add copybuttons to all of our code cells + const COPYBUTTON_SELECTOR = 'div.highlight pre'; + const codeCells = document.querySelectorAll(COPYBUTTON_SELECTOR) + codeCells.forEach((codeCell, index) => { + const id = codeCellId(index) + codeCell.setAttribute('id', id) + + const clipboardButton = id => + `` + codeCell.insertAdjacentHTML('afterend', clipboardButton(id)) + }) + +function escapeRegExp(string) { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string +} + +/** + * Removes excluded text from a Node. + * + * @param {Node} target Node to filter. + * @param {string} exclude CSS selector of nodes to exclude. + * @returns {DOMString} Text from `target` with text removed. + */ +function filterText(target, exclude) { + const clone = target.cloneNode(true); // clone as to not modify the live DOM + if (exclude) { + // remove excluded nodes + clone.querySelectorAll(exclude).forEach(node => node.remove()); + } + return clone.innerText; +} + +// Callback when a copy button is clicked. Will be passed the node that was clicked +// should then grab the text and replace pieces of text that shouldn't be used in output +function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { + var regexp; + var match; + + // Do we check for line continuation characters and "HERE-documents"? + var useLineCont = !!lineContinuationChar + var useHereDoc = !!hereDocDelim + + // create regexp to capture prompt and remaining line + if (isRegexp) { + regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') + } else { + regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') + } + + const outputLines = []; + var promptFound = false; + var gotLineCont = false; + var gotHereDoc = false; + const lineGotPrompt = []; + for (const line of textContent.split('\n')) { + match = line.match(regexp) + if (match || gotLineCont || gotHereDoc) { + promptFound = regexp.test(line) + lineGotPrompt.push(promptFound) + if (removePrompts && promptFound) { + outputLines.push(match[2]) + } else { + outputLines.push(line) + } + gotLineCont = line.endsWith(lineContinuationChar) & useLineCont + if (line.includes(hereDocDelim) & useHereDoc) + gotHereDoc = !gotHereDoc + } else if (!onlyCopyPromptLines) { + outputLines.push(line) + } else if (copyEmptyLines && line.trim() === '') { + outputLines.push(line) + } + } + + // If no lines with the prompt were found then just use original lines + if (lineGotPrompt.some(v => v === true)) { + textContent = outputLines.join('\n'); + } + + // Remove a trailing newline to avoid auto-running when pasting + if (textContent.endsWith("\n")) { + textContent = textContent.slice(0, -1) + } + return textContent +} + + +var copyTargetText = (trigger) => { + var target = document.querySelector(trigger.attributes['data-clipboard-target'].value); + + // get filtered text + let exclude = '.linenos'; + + let text = filterText(target, exclude); + return formatCopyText(text, '', false, true, true, true, '', '') +} + + // Initialize with a callback so we can modify the text before copy + const clipboard = new ClipboardJS('.copybtn', {text: copyTargetText}) + + // Update UI with error/success messages + clipboard.on('success', event => { + clearSelection() + temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_success']) + temporarilyChangeIcon(event.trigger) + }) + + clipboard.on('error', event => { + temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_failure']) + }) +} + +runWhenDOMLoaded(addCopyButtonToCodeCells) \ No newline at end of file diff --git a/_static/copybutton_funcs.js b/_static/copybutton_funcs.js new file mode 100644 index 00000000..dbe1aaad --- /dev/null +++ b/_static/copybutton_funcs.js @@ -0,0 +1,73 @@ +function escapeRegExp(string) { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string +} + +/** + * Removes excluded text from a Node. + * + * @param {Node} target Node to filter. + * @param {string} exclude CSS selector of nodes to exclude. + * @returns {DOMString} Text from `target` with text removed. + */ +export function filterText(target, exclude) { + const clone = target.cloneNode(true); // clone as to not modify the live DOM + if (exclude) { + // remove excluded nodes + clone.querySelectorAll(exclude).forEach(node => node.remove()); + } + return clone.innerText; +} + +// Callback when a copy button is clicked. Will be passed the node that was clicked +// should then grab the text and replace pieces of text that shouldn't be used in output +export function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { + var regexp; + var match; + + // Do we check for line continuation characters and "HERE-documents"? + var useLineCont = !!lineContinuationChar + var useHereDoc = !!hereDocDelim + + // create regexp to capture prompt and remaining line + if (isRegexp) { + regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') + } else { + regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') + } + + const outputLines = []; + var promptFound = false; + var gotLineCont = false; + var gotHereDoc = false; + const lineGotPrompt = []; + for (const line of textContent.split('\n')) { + match = line.match(regexp) + if (match || gotLineCont || gotHereDoc) { + promptFound = regexp.test(line) + lineGotPrompt.push(promptFound) + if (removePrompts && promptFound) { + outputLines.push(match[2]) + } else { + outputLines.push(line) + } + gotLineCont = line.endsWith(lineContinuationChar) & useLineCont + if (line.includes(hereDocDelim) & useHereDoc) + gotHereDoc = !gotHereDoc + } else if (!onlyCopyPromptLines) { + outputLines.push(line) + } else if (copyEmptyLines && line.trim() === '') { + outputLines.push(line) + } + } + + // If no lines with the prompt were found then just use original lines + if (lineGotPrompt.some(v => v === true)) { + textContent = outputLines.join('\n'); + } + + // Remove a trailing newline to avoid auto-running when pasting + if (textContent.endsWith("\n")) { + textContent = textContent.slice(0, -1) + } + return textContent +} diff --git a/_static/design-tabs.js b/_static/design-tabs.js new file mode 100644 index 00000000..b25bd6a4 --- /dev/null +++ b/_static/design-tabs.js @@ -0,0 +1,101 @@ +// @ts-check + +// Extra JS capability for selected tabs to be synced +// The selection is stored in local storage so that it persists across page loads. + +/** + * @type {Record