From 96adaec106b182b314b08bb8b85fb24ea72547e7 Mon Sep 17 00:00:00 2001 From: Andy Salnikov Date: Tue, 11 Jul 2023 15:25:50 -0700 Subject: [PATCH] Add `RegistryShim` class. This new class implements `Registry` interface and forwards all methods to other objects. Everything is forwarded to `Butler._registry` now, but this will change as we start refactoring ButlerRegistry implementation. --- python/lsst/daf/butler/_butler.py | 5 +- python/lsst/daf/butler/_registry_shim.py | 394 +++++++++++++++++++++++ tests/test_butler.py | 4 +- 3 files changed, 400 insertions(+), 3 deletions(-) create mode 100644 python/lsst/daf/butler/_registry_shim.py diff --git a/python/lsst/daf/butler/_butler.py b/python/lsst/daf/butler/_butler.py index a50298f751..7bc77e2dfd 100644 --- a/python/lsst/daf/butler/_butler.py +++ b/python/lsst/daf/butler/_butler.py @@ -51,6 +51,7 @@ from ._dataset_existence import DatasetExistence from ._deferredDatasetHandle import DeferredDatasetHandle from ._limited_butler import LimitedButler +from ._registry_shim import RegistryShim from .core import ( Config, ConfigSubset, @@ -246,6 +247,8 @@ def __init__( if "run" in self._config or "collection" in self._config: raise ValueError("Passing a run or collection via configuration is no longer supported.") + self._registry_shim = RegistryShim(self) + GENERATION: ClassVar[int] = 3 """This is a Generation 3 Butler. @@ -2636,7 +2639,7 @@ def registry(self) -> Registry: are accessible only via `Registry` methods. Eventually these methods will be replaced by equivalent `Butler` methods. """ - return self._registry + return self._registry_shim @property def dimensions(self) -> DimensionUniverse: diff --git a/python/lsst/daf/butler/_registry_shim.py b/python/lsst/daf/butler/_registry_shim.py new file mode 100644 index 0000000000..81a184520d --- /dev/null +++ b/python/lsst/daf/butler/_registry_shim.py @@ -0,0 +1,394 @@ +# This file is part of daf_butler. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (http://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +from __future__ import annotations + +__all__ = ("Registry",) + +import contextlib +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any + +from .core import ( + DataCoordinate, + DataId, + DatasetAssociation, + DatasetId, + DatasetIdGenEnum, + DatasetRef, + DatasetType, + Dimension, + DimensionElement, + DimensionGraph, + DimensionRecord, + DimensionUniverse, + NameLookupMapping, + Timespan, +) +from .registry import Registry +from .registry._collection_summary import CollectionSummary +from .registry._collectionType import CollectionType +from .registry._defaults import RegistryDefaults +from .registry.queries import DataCoordinateQueryResults, DatasetQueryResults, DimensionRecordQueryResults + +if TYPE_CHECKING: + from ._butler import Butler + from .registry._registry import CollectionArgType + from .registry.interfaces import ObsCoreTableManager + + +class RegistryShim(Registry): + """Implementation of `Registry` interface exposed to clients by `Butler`. + + Parameters + ---------- + butler : `Butler` + Data butler instance. + + Notes + ----- + This shim implementation of `Registry` forwards all methods to an actual + Registry instance which is internal to Butler or to Butler methods. Its + purpose is to provide a stable interface to many client-visible operations + while we perform re-structuring of Registry and Butler implementations. + """ + + def __init__(self, butler: Butler): + self._butler = butler + self._registry = butler._registry + + def isWriteable(self) -> bool: + # Docstring inherited from a base class. + return self._registry.isWriteable() + + @property + def dimensions(self) -> DimensionUniverse: + # Docstring inherited from a base class. + return self._registry.dimensions + + @property + def defaults(self) -> RegistryDefaults: + # Docstring inherited from a base class. + + # Note that we do not allow changing defaults via this class + return self._registry.defaults + + @defaults.setter + def defaults(self, value: RegistryDefaults) -> None: + # Docstring inherited from a base class. + self._registry.defaults = value + + def refresh(self) -> None: + # Docstring inherited from a base class. + self._registry.refresh() + + @contextlib.contextmanager + def transaction(self, *, savepoint: bool = False) -> Iterator[None]: + # Docstring inherited from a base class. + with self._registry.transaction(savepoint=savepoint): + yield + + def resetConnectionPool(self) -> None: + # Docstring inherited from a base class. + self._registry.resetConnectionPool() + + def registerCollection( + self, name: str, type: CollectionType = CollectionType.TAGGED, doc: str | None = None + ) -> bool: + # Docstring inherited from a base class. + return self._registry.registerCollection(name, type, doc) + + def getCollectionType(self, name: str) -> CollectionType: + # Docstring inherited from a base class. + return self._registry.getCollectionType(name) + + def registerRun(self, name: str, doc: str | None = None) -> bool: + # Docstring inherited from a base class. + return self._registry.registerRun(name, doc) + + def removeCollection(self, name: str) -> None: + # Docstring inherited from a base class. + self._registry.removeCollection(name) + + def getCollectionChain(self, parent: str) -> Sequence[str]: + # Docstring inherited from a base class. + return self._registry.getCollectionChain(parent) + + def setCollectionChain(self, parent: str, children: Any, *, flatten: bool = False) -> None: + # Docstring inherited from a base class. + self._registry.setCollectionChain(parent, children, flatten=flatten) + + def getCollectionParentChains(self, collection: str) -> set[str]: + # Docstring inherited from a base class. + return self._registry.getCollectionParentChains(collection) + + def getCollectionDocumentation(self, collection: str) -> str | None: + # Docstring inherited from a base class. + return self._registry.getCollectionDocumentation(collection) + + def setCollectionDocumentation(self, collection: str, doc: str | None) -> None: + # Docstring inherited from a base class. + self._registry.setCollectionDocumentation(collection, doc) + + def getCollectionSummary(self, collection: str) -> CollectionSummary: + # Docstring inherited from a base class. + return self._registry.getCollectionSummary(collection) + + def registerDatasetType(self, datasetType: DatasetType) -> bool: + # Docstring inherited from a base class. + return self._registry.registerDatasetType(datasetType) + + def removeDatasetType(self, name: str | tuple[str, ...]) -> None: + # Docstring inherited from a base class. + self._registry.removeDatasetType(name) + + def getDatasetType(self, name: str) -> DatasetType: + # Docstring inherited from a base class. + return self._registry.getDatasetType(name) + + def supportsIdGenerationMode(self, mode: DatasetIdGenEnum) -> bool: + # Docstring inherited from a base class. + return self._registry.supportsIdGenerationMode(mode) + + def findDataset( + self, + datasetType: DatasetType | str, + dataId: DataId | None = None, + *, + collections: CollectionArgType | None = None, + timespan: Timespan | None = None, + **kwargs: Any, + ) -> DatasetRef | None: + # Docstring inherited from a base class. + return self._registry.findDataset( + datasetType, dataId, collections=collections, timespan=timespan, **kwargs + ) + + def insertDatasets( + self, + datasetType: DatasetType | str, + dataIds: Iterable[DataId], + run: str | None = None, + expand: bool = True, + idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE, + ) -> list[DatasetRef]: + # Docstring inherited from a base class. + return self._registry.insertDatasets(datasetType, dataIds, run, expand, idGenerationMode) + + def _importDatasets(self, datasets: Iterable[DatasetRef], expand: bool = True) -> list[DatasetRef]: + # Docstring inherited from a base class. + return self._registry._importDatasets(datasets, expand) + + def getDataset(self, id: DatasetId) -> DatasetRef | None: + # Docstring inherited from a base class. + return self._registry.getDataset(id) + + def removeDatasets(self, refs: Iterable[DatasetRef]) -> None: + # Docstring inherited from a base class. + self._registry.removeDatasets(refs) + + def associate(self, collection: str, refs: Iterable[DatasetRef]) -> None: + # Docstring inherited from a base class. + self._registry.associate(collection, refs) + + def disassociate(self, collection: str, refs: Iterable[DatasetRef]) -> None: + # Docstring inherited from a base class. + self._registry.disassociate(collection, refs) + + def certify(self, collection: str, refs: Iterable[DatasetRef], timespan: Timespan) -> None: + # Docstring inherited from a base class. + self._registry.certify(collection, refs, timespan) + + def decertify( + self, + collection: str, + datasetType: str | DatasetType, + timespan: Timespan, + *, + dataIds: Iterable[DataId] | None = None, + ) -> None: + # Docstring inherited from a base class. + self._registry.decertify(collection, datasetType, timespan, dataIds=dataIds) + + def getDatasetLocations(self, ref: DatasetRef) -> Iterable[str]: + # Docstring inherited from a base class. + return self._registry.getDatasetLocations(ref) + + def expandDataId( + self, + dataId: DataId | None = None, + *, + graph: DimensionGraph | None = None, + records: NameLookupMapping[DimensionElement, DimensionRecord | None] | None = None, + withDefaults: bool = True, + **kwargs: Any, + ) -> DataCoordinate: + # Docstring inherited from a base class. + return self._registry.expandDataId( + dataId, graph=graph, records=records, withDefaults=withDefaults, **kwargs + ) + + def insertDimensionData( + self, + element: DimensionElement | str, + *data: Mapping[str, Any] | DimensionRecord, + conform: bool = True, + replace: bool = False, + skip_existing: bool = False, + ) -> None: + # Docstring inherited from a base class. + self._registry.insertDimensionData( + element, *data, conform=conform, replace=replace, skip_existing=skip_existing + ) + + def syncDimensionData( + self, + element: DimensionElement | str, + row: Mapping[str, Any] | DimensionRecord, + conform: bool = True, + update: bool = False, + ) -> bool | dict[str, Any]: + # Docstring inherited from a base class. + return self._registry.syncDimensionData(element, row, conform, update) + + def queryDatasetTypes( + self, + expression: Any = ..., + *, + components: bool | None = None, + missing: list[str] | None = None, + ) -> Iterable[DatasetType]: + # Docstring inherited from a base class. + return self._registry.queryDatasetTypes(expression, components=components, missing=missing) + + def queryCollections( + self, + expression: Any = ..., + datasetType: DatasetType | None = None, + collectionTypes: Iterable[CollectionType] | CollectionType = CollectionType.all(), + flattenChains: bool = False, + includeChains: bool | None = None, + ) -> Sequence[str]: + # Docstring inherited from a base class. + return self._registry.queryCollections( + expression, datasetType, collectionTypes, flattenChains, includeChains + ) + + def queryDatasets( + self, + datasetType: Any, + *, + collections: CollectionArgType | None = None, + dimensions: Iterable[Dimension | str] | None = None, + dataId: DataId | None = None, + where: str = "", + findFirst: bool = False, + components: bool | None = None, + bind: Mapping[str, Any] | None = None, + check: bool = True, + **kwargs: Any, + ) -> DatasetQueryResults: + # Docstring inherited from a base class. + return self._registry.queryDatasets( + datasetType, + collections=collections, + dimensions=dimensions, + dataId=dataId, + where=where, + findFirst=findFirst, + components=components, + bind=bind, + check=check, + **kwargs, + ) + + def queryDataIds( + self, + dimensions: Iterable[Dimension | str] | Dimension | str, + *, + dataId: DataId | None = None, + datasets: Any = None, + collections: CollectionArgType | None = None, + where: str = "", + components: bool | None = None, + bind: Mapping[str, Any] | None = None, + check: bool = True, + **kwargs: Any, + ) -> DataCoordinateQueryResults: + # Docstring inherited from a base class. + return self._registry.queryDataIds( + dimensions, + dataId=dataId, + datasets=datasets, + collections=collections, + where=where, + components=components, + bind=bind, + check=check, + **kwargs, + ) + + def queryDimensionRecords( + self, + element: DimensionElement | str, + *, + dataId: DataId | None = None, + datasets: Any = None, + collections: CollectionArgType | None = None, + where: str = "", + components: bool | None = None, + bind: Mapping[str, Any] | None = None, + check: bool = True, + **kwargs: Any, + ) -> DimensionRecordQueryResults: + # Docstring inherited from a base class. + return self._registry.queryDimensionRecords( + element, + dataId=dataId, + datasets=datasets, + collections=collections, + where=where, + components=components, + bind=bind, + check=check, + **kwargs, + ) + + def queryDatasetAssociations( + self, + datasetType: str | DatasetType, + collections: CollectionArgType | None = ..., + *, + collectionTypes: Iterable[CollectionType] = CollectionType.all(), + flattenChains: bool = False, + ) -> Iterator[DatasetAssociation]: + # Docstring inherited from a base class. + return self._registry.queryDatasetAssociations( + datasetType, + collections, + collectionTypes=collectionTypes, + flattenChains=flattenChains, + ) + + @property + def obsCoreTableManager(self) -> ObsCoreTableManager | None: + # Docstring inherited from a base class. + return self._registry.obsCoreTableManager diff --git a/tests/test_butler.py b/tests/test_butler.py index 605f986227..307cad7573 100644 --- a/tests/test_butler.py +++ b/tests/test_butler.py @@ -1679,8 +1679,8 @@ def testPytypeCoercion(self) -> None: # Now need to hack the registry dataset type definition. # There is no API for this. - assert isinstance(butler.registry, SqlRegistry) - manager = butler.registry._managers.datasets + assert isinstance(butler._registry, SqlRegistry) + manager = butler._registry._managers.datasets assert hasattr(manager, "_db") and hasattr(manager, "_static") manager._db.update( manager._static.dataset_type,