Improve type checking fot generic datastore

python/lsst/daf/butler/datastores/fileDatastore.py

@@ -126,7 +126,7 @@ class DatastoreFileGetInformation:
     """The `StorageClass` of the dataset being read."""
 
 
-class FileDatastore(GenericBaseDatastore):
+class FileDatastore(GenericBaseDatastore[StoredFileInfo]):
     """Generic Datastore for file-based implementations.
 
     Should always be sub-classed since key abstract methods are missing.

python/lsst/daf/butler/datastores/genericDatastore.py

@@ -27,19 +27,21 @@
 
 import logging
 from abc import abstractmethod
-from collections.abc import Iterable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-from lsst.daf.butler import DatasetTypeNotSupportedError, Datastore
+from lsst.daf.butler import DatasetTypeNotSupportedError, Datastore, StoredDatastoreItemInfo
 from lsst.daf.butler.registry.interfaces import DatastoreRegistryBridge
 
 if TYPE_CHECKING:
-    from lsst.daf.butler import DatasetRef, StorageClass, StoredDatastoreItemInfo
+    from lsst.daf.butler import DatasetRef, StorageClass
 
 log = logging.getLogger(__name__)
 
+_InfoType = TypeVar("_InfoType", bound=StoredDatastoreItemInfo)
 
-class GenericBaseDatastore(Datastore):
+
+class GenericBaseDatastore(Datastore, Generic[_InfoType]):
     """Methods useful for most implementations of a `Datastore`.
 
     Should always be sub-classed since key abstract methods are missing.
@@ -54,7 +56,7 @@ def bridge(self) -> DatastoreRegistryBridge:
         raise NotImplementedError()
 
     @abstractmethod
-    def addStoredItemInfo(self, refs: Iterable[DatasetRef], infos: Iterable[Any]) -> None:
+    def addStoredItemInfo(self, refs: Iterable[DatasetRef], infos: Iterable[_InfoType]) -> None:
         """Record internal storage information associated with one or more
         datasets.
 
@@ -68,7 +70,7 @@ def addStoredItemInfo(self, refs: Iterable[DatasetRef], infos: Iterable[Any]) ->
         raise NotImplementedError()
 
     @abstractmethod
-    def getStoredItemsInfo(self, ref: DatasetRef) -> Sequence[Any]:
+    def getStoredItemsInfo(self, ref: DatasetRef) -> Iterable[_InfoType]:
         """Retrieve information associated with files stored in this
         `Datastore` associated with this dataset ref.
 
@@ -79,7 +81,7 @@ def getStoredItemsInfo(self, ref: DatasetRef) -> Sequence[Any]:
 
         Returns
         -------
-        items : `list` [`StoredDatastoreItemInfo`]
+        items : `~collections.abc.Iterable` [`StoredDatastoreItemInfo`]
             Stored information about the files and associated formatters
             associated with this dataset. Only one file will be returned
             if the dataset has not been disassembled. Can return an empty
@@ -98,7 +100,7 @@ def removeStoredItemInfo(self, ref: DatasetRef) -> None:
         """
         raise NotImplementedError()
 
-    def _register_datasets(self, refsAndInfos: Iterable[tuple[DatasetRef, StoredDatastoreItemInfo]]) -> None:
+    def _register_datasets(self, refsAndInfos: Iterable[tuple[DatasetRef, _InfoType]]) -> None:
         """Update registry to indicate that one or more datasets have been
         stored.
 
@@ -125,11 +127,11 @@ def _register_datasets(self, refsAndInfos: Iterable[tuple[DatasetRef, StoredData
 
     def _post_process_get(
         self,
-        inMemoryDataset: Any,
+        inMemoryDataset: object,
         readStorageClass: StorageClass,
         assemblerParams: Mapping[str, Any] | None = None,
         isComponent: bool = False,
-    ) -> Any:
+    ) -> object:
         """Given the Python object read from the datastore, manipulate
         it based on the supplied parameters and ensure the Python
         type is correct.
@@ -145,6 +147,11 @@ def _post_process_get(
             Parameters to pass to the assembler.  Can be `None`.
         isComponent : `bool`, optional
             If this is a component, allow the inMemoryDataset to be `None`.
+
+        Returns
+        -------
+        dataset : `object`
+            In-memory dataset, potentially converted to expected type.
         """
         # Process any left over parameters
         if assemblerParams:
@@ -166,7 +173,7 @@ def _post_process_get(
 
         return inMemoryDataset
 
-    def _validate_put_parameters(self, inMemoryDataset: Any, ref: DatasetRef) -> None:
+    def _validate_put_parameters(self, inMemoryDataset: object, ref: DatasetRef) -> None:
         """Validate the supplied arguments for put.
 
         Parameters

python/lsst/daf/butler/datastores/inMemoryDatastore.py

@@ -74,7 +74,7 @@ class StoredMemoryItemInfo(StoredDatastoreItemInfo):
     """
 
 
-class InMemoryDatastore(GenericBaseDatastore):
+class InMemoryDatastore(GenericBaseDatastore[StoredMemoryItemInfo]):
     """Basic Datastore for writing to an in memory cache.
 
     This datastore is ephemeral in that the contents of the datastore