Skip to content

Commit

Permalink
Merge pull request #1052 from lsst/tickets/DM-45726
Browse files Browse the repository at this point in the history
DM-45726: Pass Butler to RepoExportContext rather than Registry
  • Loading branch information
timj authored Aug 13, 2024
2 parents a0830e1 + aef9274 commit 20b5ab1
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 23 deletions.
4 changes: 1 addition & 3 deletions python/lsst/daf/butler/direct_butler/_direct_butler.py
Original file line number Diff line number Diff line change
Expand Up @@ -1607,9 +1607,7 @@ def export(
with open(filename, "w") as stream:
backend = BackendClass(stream, universe=self.dimensions)
try:
helper = RepoExportContext(
self._registry, self._datastore, backend=backend, directory=directory, transfer=transfer
)
helper = RepoExportContext(self, backend=backend, directory=directory, transfer=transfer)
with self._caching_context():
yield helper
except BaseException:
Expand Down
35 changes: 15 additions & 20 deletions python/lsst/daf/butler/transfers/_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,15 +37,14 @@
from .._dataset_ref import DatasetId, DatasetRef
from .._dataset_type import DatasetType
from .._file_dataset import FileDataset
from ..datastore import Datastore
from ..dimensions import DataCoordinate, DimensionElement, DimensionRecord
from ..registry import CollectionType
from ..registry.interfaces import ChainedCollectionRecord, CollectionRecord

if TYPE_CHECKING:
from lsst.resources import ResourcePathExpression

from ..registry.sql_registry import SqlRegistry
from ..direct_butler import DirectButler
from ._interfaces import RepoExportBackend


Expand All @@ -61,10 +60,8 @@ class RepoExportContext:
Parameters
----------
registry : `SqlRegistry`
Registry to export from.
datastore : `Datastore`
Datastore to export from.
butler : `lsst.daf.butler.direct_butler.DirectButler`
Butler to export from.
backend : `RepoExportBackend`
Implementation class for a particular export file format.
directory : `~lsst.resources.ResourcePathExpression`, optional
Expand All @@ -76,15 +73,13 @@ class RepoExportContext:

def __init__(
self,
registry: SqlRegistry,
datastore: Datastore,
butler: DirectButler, # Requires butler._registry to work for now.
backend: RepoExportBackend,
*,
directory: ResourcePathExpression | None = None,
transfer: str | None = None,
):
self._registry = registry
self._datastore = datastore
self._butler = butler
self._backend = backend
self._directory = directory
self._transfer = transfer
Expand Down Expand Up @@ -118,7 +113,7 @@ def saveCollection(self, name: str) -> None:
export its child collections; these must be explicitly exported or
already be present in the repository they are being imported into.
"""
self._collections[name] = self._registry.get_collection_record(name)
self._collections[name] = self._butler._registry.get_collection_record(name)

def saveDimensionData(
self, element: str | DimensionElement, records: Iterable[dict | DimensionRecord]
Expand All @@ -136,7 +131,7 @@ def saveDimensionData(
`dict` instances.
"""
if not isinstance(element, DimensionElement):
element = self._registry.dimensions[element]
element = self._butler.dimensions[element]
for record in records:
if not isinstance(record, DimensionRecord):
record = element.RecordClass(**record)
Expand Down Expand Up @@ -170,13 +165,13 @@ def saveDataIds(
standardized_elements: Set[DimensionElement]
if elements is None:
standardized_elements = frozenset(
element for element in self._registry.dimensions.elements if element.has_own_table
element for element in self._butler.dimensions.elements if element.has_own_table
)
else:
standardized_elements = set()
for element in elements:
if not isinstance(element, DimensionElement):
element = self._registry.dimensions[element]
element = self._butler.dimensions[element]
if element.has_own_table:
standardized_elements.add(element)
for dataId in dataIds:
Expand All @@ -185,7 +180,7 @@ def saveDataIds(
# if the data ID is already expanded, and DM-26692 will add (or at
# least start to add / unblock) query functionality that should
# let us speed this up internally as well.
dataId = self._registry.expandDataId(dataId)
dataId = self._butler.registry.expandDataId(dataId)
for element_name in dataId.dimensions.elements:
record = dataId.records[element_name]
if record is not None and record.definition in standardized_elements:
Expand Down Expand Up @@ -243,7 +238,7 @@ def saveDatasets(
refs_to_export[dataset_id] = ref
# Do a vectorized datastore export, which might be a lot faster than
# one-by-one.
exports = self._datastore.export(
exports = self._butler._datastore.export(
refs_to_export.values(),
directory=self._directory,
transfer=self._transfer,
Expand All @@ -267,15 +262,15 @@ def _finish(self) -> None:
For use by `Butler.export` only.
"""
for element in self._registry.dimensions.sorted(self._records.keys()):
for element in self._butler.dimensions.sorted(self._records.keys()):
# To make export deterministic sort the DataCoordinate instances.
r = self._records[element]
self._backend.saveDimensionData(element, *[r[dataId] for dataId in sorted(r.keys())])
for datasetsByRun in self._datasets.values():
for run in datasetsByRun:
self._collections[run] = self._registry.get_collection_record(run)
self._collections[run] = self._butler._registry.get_collection_record(run)
for collectionName in self._computeSortedCollections():
doc = self._registry.getCollectionDocumentation(collectionName)
doc = self._butler.registry.getCollectionDocumentation(collectionName)
self._backend.saveCollection(self._collections[collectionName], doc)
# Sort the dataset types and runs before exporting to ensure
# reproducible order in export file.
Expand Down Expand Up @@ -363,7 +358,7 @@ def _computeDatasetAssociations(self) -> dict[str, list[DatasetAssociation]]:
collectionTypes = {CollectionType.TAGGED}
if datasetType.isCalibration():
collectionTypes.add(CollectionType.CALIBRATION)
associationIter = self._registry.queryDatasetAssociations(
associationIter = self._butler.registry.queryDatasetAssociations(
datasetType,
collections=self._collections.keys(),
collectionTypes=collectionTypes,
Expand Down

0 comments on commit 20b5ab1

Please sign in to comment.