Merge pull request #19 from scipp/some-typing-improvements

Some typing improvements

pyproject.toml

@@ -14,3 +14,7 @@ filterwarnings = [
   "error",
   "ignore::UserWarning",
 ]
+
+[tool.mypy]
+mypy_path = "src"
+ignore_missing_imports = true

setup.cfg

@@ -2,13 +2,19 @@
 name = scippnexus
 author = Scipp contributors (https://github.com/scipp)
 description = h5py-like utility for NeXus files based with seamless scipp integration
+license = BSD
+license_file = LICENSE
 long_description = file: README.md
 long_description_content_type = text/markdown
 url = https://scipp.github.io/scippnexus
 project_urls =
     Bug Tracker = https://github.com/scipp/scippnexus/issues
 classifiers =
     Programming Language :: Python :: 3
+    Programming Language :: Python :: 3 :: Only
+    Programming Language :: Python :: 3.8
+    Programming Language :: Python :: 3.9
+    Programming Language :: Python :: 3.10
     License :: OSI Approved :: BSD License
     Operating System :: OS Independent
 
@@ -26,6 +32,9 @@ include_package_data = True
 [options.packages.find]
 where = src
 
+[options.package_data]
+scippnexus = py.typed
+
 [flake8]
 # See https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#line-length
 max-line-length = 88

src/scippnexus/_common.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2022 Scipp contributors (https://github.com/scipp)
 # @author Simon Heybrock
-from typing import Union, List
+from typing import Dict, Union, List
 import scipp as sc
 import numpy as np
 from .typing import ScippIndex
@@ -48,7 +48,8 @@ def convert_time_to_datetime64(
         dtype=sc.DType.int64, unit=unit, copy=False)
 
 
-def _to_canonical_select(dims: List[str], select: ScippIndex) -> ScippIndex:
+def _to_canonical_select(dims: List[str],
+                         select: ScippIndex) -> Dict[str, Union[int, slice]]:
     """Return selection as dict with explicit dim labels"""
     def check_1d():
         if len(dims) != 1:
@@ -76,7 +77,7 @@ def check_1d():
     return select
 
 
-def to_plain_index(dims: List[str], select: ScippIndex) -> Union[int, tuple]:
+def to_plain_index(dims: List[str], select: ScippIndex) -> Union[int, slice, tuple]:
     """
     Given a valid "scipp" index 'select', return an equivalent plain numpy-style index.
     """

src/scippnexus/nxdata.py

@@ -159,7 +159,7 @@ def _getitem(self, select: ScippIndex) -> sc.DataArray:
                 continue
             try:
                 sel = to_child_select(self.dims, field.dims, select)
-                coord = self[name][sel]
+                coord: sc.Variable = self[name][sel]
                 for suffix in self._error_suffixes:
                     if f'{name}{suffix}' in self:
                         if coord.variances is not None:

src/scippnexus/nxdetector.py

@@ -60,7 +60,7 @@ def unit(self):
         self._nxevent_data.unit
 
     def __getitem__(self, select: ScippIndex) -> sc.DataArray:
-        event_data = self._nxevent_data[self._event_select]
+        event_data: sc.DataArray = self._nxevent_data[self._event_select]
         if self._detector_number is None:
             if select not in (Ellipsis, tuple(), slice(None)):
                 raise NexusStructureError(

src/scippnexus/nxevent_data.py

@@ -37,7 +37,7 @@ def _getitem(self, select: ScippIndex) -> sc.DataArray:
         self._check_for_missing_fields()
         index = to_plain_index([_pulse_dimension], select)
 
-        max_index = self["event_index"].shape[0]
+        max_index = self.shape[0]
         single = False
         if index is Ellipsis or index == tuple():
             last_loaded = False

src/scippnexus/nxobject.py

@@ -8,7 +8,7 @@
 import dateutil.parser
 from enum import Enum, auto
 import functools
-from typing import List, Union, NoReturn, Any, Dict, Tuple, Protocol
+from typing import List, Union, Any, Dict, Tuple, Protocol
 import numpy as np
 import scipp as sc
 import h5py
@@ -291,11 +291,12 @@ def _get_child(
             da.coords['depends_on'] = t if isinstance(t, sc.Variable) else sc.scalar(t)
         return da
 
-    def __getitem__(self,
-                    name: NXobjectIndex) -> Union['NXobject', Field, sc.DataArray]:
+    def __getitem__(
+            self,
+            name: NXobjectIndex) -> Union['NXobject', Field, sc.DataArray, sc.Dataset]:
         return self._get_child(name, use_field_dims=True)
 
-    def _getitem(self, index: ScippIndex) -> NoReturn:
+    def _getitem(self, index: ScippIndex) -> Union[sc.DataArray, sc.Dataset]:
         raise NotImplementedError(f'Loading {self.nx_class} is not supported.')
 
     def _get_field_dims(self, name: str) -> Union[None, List[str]]:

src/scippnexus/typing.py

@@ -1,24 +1,71 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2022 Scipp contributors (https://github.com/scipp)
 # @author Simon Heybrock
-from typing import Protocol, Union, Tuple, Dict, List, Callable
+from __future__ import annotations
+from typing import Any, Protocol, Union, Tuple, Dict, List, Callable
+from typing import TYPE_CHECKING
+
+
+class H5Base(Protocol):
+    @property
+    def attrs(self) -> List[int]:
+        """Attributes of dataset or group"""
+
+    @property
+    def name(self) -> str:
+        """Name of dataset or group"""
+
+    @property
+    def file(self) -> List[int]:
+        """File of dataset or group"""
+
+    @property
+    def parent(self) -> H5Group:
+        """Parent of dataset or group"""
 
 
 # TODO Define more required methods
-class H5Dataset(Protocol):
+class H5Dataset(H5Base, Protocol):
     """h5py.Dataset-like"""
     @property
     def shape(self) -> List[int]:
         """Shape of a dataset"""
 
+    @property
+    def dtype(self) -> List[int]:
+        """dtype of a dataset"""
+
+    def read_direct(self, array) -> None:
+        """Read dataset into given buffer"""
+
 
-class H5Group(Protocol):
+class H5Group(H5Base, Protocol):
     """h5py.Group-like"""
+    def __getitem__(self, index: Union[str, Any]) -> Union[H5Dataset, H5Group]:
+        """Keys in the group"""
+
+    def keys(self) -> List[str]:
+        """Keys in the group"""
+
+    def create_dataset(self) -> H5Dataset:
+        """Create a dataset"""
+
+    def create_group(self) -> H5Group:
+        """Create a group"""
+
     def visititems(self, func: Callable) -> None:
-        """"""
+        """Apply callable to all items, recursively"""
+
+
+if TYPE_CHECKING:
+    from enum import Enum
 
+    class ellipsis(Enum):
+        Ellipsis = "..."
+else:
+    ellipsis = type(Ellipsis)
 
 # Note that scipp does not support dicts yet, but this HDF5 code does, to
 # allow for loading blocks of 2d (or higher) data efficiently.
-ScippIndex = Union[type(Ellipsis), int, slice, Tuple[str, Union[int, slice]],
+ScippIndex = Union[ellipsis, int, tuple, slice, Tuple[str, Union[int, slice]],
                    Dict[str, Union[int, slice]]]