All seems to pass

python/lsst/daf/butler/core/datasets/ref.py

@@ -24,7 +24,6 @@
 
 import hashlib
 from typing import (
-    TYPE_CHECKING,
     Any,
     Dict,
     Iterable,
@@ -42,9 +41,6 @@
 from ..utils import immutable, NamedKeyDict
 from .type import DatasetType
 
-if TYPE_CHECKING:
-    from ...registry.interfaces import FakeDatasetRef
-
 
 class AmbiguousDatasetError(Exception):
     """Exception raised when a `DatasetRef` is not resolved (has no ID, run, or
@@ -150,7 +146,7 @@ def __new__(cls, datasetType: DatasetType, dataId: DataCoordinate, *,
             self._hash = hash
         return self
 
-    def __eq__(self, other: DatasetRef):
+    def __eq__(self, other: Any) -> bool:
         try:
             return (self.datasetType, self.dataId, self.id) == (other.datasetType, other.dataId, other.id)
         except AttributeError:
@@ -301,7 +297,7 @@ def isComposite(self) -> bool:
         """
         return self.datasetType.isComposite()
 
-    def _lookupNames(self) -> Tuple[LookupKey]:
+    def _lookupNames(self) -> Tuple[LookupKey, ...]:
         """Name keys to use when looking up this DatasetRef in a configuration.
 
         The names are returned in order of priority.
@@ -316,16 +312,45 @@ def _lookupNames(self) -> Tuple[LookupKey]:
         """
         # Special case the instrument Dimension since we allow configs
         # to include the instrument name in the hierarchy.
-        names = self.datasetType._lookupNames()
+        names: Tuple[LookupKey, ...] = self.datasetType._lookupNames()
 
         if "instrument" in self.dataId:
             names = tuple(n.clone(dataId={"instrument": self.dataId["instrument"]})
                           for n in names) + names
 
         return names
 
+    def allRefs(self, parents: bool = True) -> Iterator[DatasetRef]:
+        """Return all the nested component `DatasetRef` and optionally the
+        parent.
+
+        Parameters
+        ----------
+        parents : `bool`, optional
+            If `True` (default) include the given dataset in the output
+            iterable.  If `False`, include only its components.  This does
+            not propagate recursively - only the outermost level of parents
+            is ignored if ``parents`` is `False`.
+
+        Yields
+        ------
+        ref : `DatasetRef`
+            Itself (only if ``parent`` is `True`) or one of its (recursive)
+            children.
+
+        Notes
+        -----
+        If ``parents`` is `True`, components are guaranteed to be yielded
+        before their parents.
+        """
+        if self.components is None:
+            raise AmbiguousDatasetError(f"Unresolved ref {self} cannot be flattened.")
+        yield from DatasetRef.flatten(self.components.values(), parents=True)
+        if parents:
+            yield self
+
     @staticmethod
-    def flatten(refs: Iterable[Union[FakeDatasetRef, DatasetRef]], *,
+    def flatten(refs: Iterable[DatasetRef], *,
                 parents: bool = True) -> Iterator[DatasetRef]:
         """Recursively transform an iterable over `DatasetRef` to include
         nested component `DatasetRef` instances.
@@ -345,19 +370,16 @@ def flatten(refs: Iterable[Union[FakeDatasetRef, DatasetRef]], *,
         ------
         ref : `DatasetRef`
             Either one of the given `DatasetRef` instances (only if ``parent``
-            is `True`) or on of its (recursive) children.
+            is `True`) or one of its (recursive) children.
 
         Notes
         -----
         If ``parents`` is `True`, components are guaranteed to be yielded
         before their parents.
         """
         for ref in refs:
-            if ref.components is None:
-                raise AmbiguousDatasetError(f"Unresolved ref {ref} passed to 'flatten'.")
-            yield from DatasetRef.flatten(ref.components.values(), parents=True)
-            if parents:
-                yield ref
+            for subref in ref.allRefs(parents):
+                yield subref
 
     @staticmethod
     def groupByType(refs: Iterable[DatasetRef], *, recursive: bool = True

python/lsst/daf/butler/core/datasets/type.py

@@ -35,6 +35,7 @@
     Mapping,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -121,6 +122,7 @@ def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimensi
             dimensions = universe.extract(dimensions)
         self._dimensions = dimensions
         assert isinstance(storageClass, (StorageClass, str))
+        self._storageClass: Optional[StorageClass]
         if isinstance(storageClass, StorageClass):
             self._storageClass = storageClass
             self._storageClassName = storageClass.name
@@ -131,7 +133,7 @@ def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimensi
     def __repr__(self) -> str:
         return "DatasetType({}, {}, {})".format(self.name, self.dimensions, self._storageClassName)
 
-    def __eq__(self, other: DatasetType) -> bool:
+    def __eq__(self, other: Any) -> bool:
         if not isinstance(other, type(self)):
             return False
         if self._name != other._name:
@@ -177,7 +179,7 @@ def storageClass(self) -> StorageClass:
         return self._storageClass
 
     @staticmethod
-    def splitDatasetTypeName(datasetTypeName: str) -> Tuple[str, str]:
+    def splitDatasetTypeName(datasetTypeName: str) -> Tuple[str, Optional[str]]:
         """Given a dataset type name, return the root name and the component
         name.
 
@@ -206,7 +208,7 @@ def splitDatasetTypeName(datasetTypeName: str) -> Tuple[str, str]:
             root, comp = root.split(".", maxsplit=1)
         return root, comp
 
-    def nameAndComponent(self):
+    def nameAndComponent(self) -> Tuple[str, Optional[str]]:
         """Return the root name of this dataset type and the component
         name (if defined).
 
@@ -309,7 +311,7 @@ def _lookupNames(self) -> Tuple[LookupKey, ...]:
             the storage class name.
         """
         rootName, componentName = self.nameAndComponent()
-        lookups = (LookupKey(name=self.name),)
+        lookups: Tuple[LookupKey, ...] = (LookupKey(name=self.name),)
         if componentName is not None:
             lookups = lookups + (LookupKey(name=rootName),)
 
@@ -319,7 +321,7 @@ def _lookupNames(self) -> Tuple[LookupKey, ...]:
 
         return lookups + self.storageClass._lookupNames()
 
-    def __reduce__(self) -> Tuple[DatasetType, Tuple[str, DimensionGraph, str]]:
+    def __reduce__(self) -> Tuple[Type[DatasetType], Tuple[str, DimensionGraph, str]]:
         """Support pickling.
 
         StorageClass instances can not normally be pickled, so we pickle

python/lsst/daf/butler/core/datastore.py

@@ -30,19 +30,22 @@
 import contextlib
 import logging
 from collections import defaultdict
-from typing import (TYPE_CHECKING,
-                    Any,
-                    Callable,
-                    ClassVar,
-                    Generator,
-                    Iterable,
-                    List,
-                    Mapping,
-                    Optional,
-                    Set,
-                    Type,
-                    Union,
-                    )
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ClassVar,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
+
 from dataclasses import dataclass
 from abc import ABCMeta, abstractmethod
 
@@ -133,7 +136,7 @@ def registerUndo(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any)
         self._log.append(self.Event(name, undoFunc, args, kwargs))
 
     @contextlib.contextmanager
-    def undoWith(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any) -> Generator:
+    def undoWith(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any) -> Iterator[None]:
         """A context manager that calls `registerUndo` if the nested operation
         does not raise an exception.
 
@@ -295,16 +298,16 @@ def __repr__(self) -> str:
         return self.name
 
     @property
-    def names(self) -> List[str]:
+    def names(self) -> Tuple[str, ...]:
         """Names associated with this datastore returned as a list.
 
         Can be different to ``name`` for a chaining datastore.
         """
         # Default implementation returns solely the name itself
-        return [self.name]
+        return (self.name, )
 
     @contextlib.contextmanager
-    def transaction(self) -> Generator:
+    def transaction(self) -> Iterator[DatastoreTransaction]:
         """Context manager supporting `Datastore` transactions.
 
         Transactions can be nested, and are to be used in combination with