Skip to content

Commit

Permalink
Add methods to update run collection name in a quantum (DM-40392)
Browse files Browse the repository at this point in the history
QuantumGraph.updateRun() method needs a better implementation and
these new methods in DatasetRef and Quantum classes are for supporting
that implementation.
  • Loading branch information
andy-slac committed Aug 18, 2023
1 parent c6c2f70 commit a900423
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 3 deletions.
43 changes: 40 additions & 3 deletions python/lsst/daf/butler/core/datasets/ref.py
Original file line number Diff line number Diff line change
Expand Up @@ -706,11 +706,48 @@ def overrideStorageClass(self, storageClass: str | StorageClass) -> DatasetRef:
A new dataset reference that is the same as the current one but
with a different storage class in the `DatasetType`.
"""
return self.replace(storage_class=storageClass)

def replace(
self,
*,
id: DatasetId | None = None,
run: str | None = None,
storage_class: str | StorageClass | None = None,
) -> DatasetRef:
"""Create a new `DatasetRef` from this one, but with some modified
attributes.
Parameters
----------
id : `DatasetId` or `None`
If not `None` then update dataset ID.
run : `str` or `None`
If not `None` then update run collection name. If ``dataset_id`` is
`None` then this will also cause new dataset ID to be generated.
storage_class : `str` or `StorageClass` or `None`.
The new storage class. If not `None`, replaces existing storage
class.
Returns
-------
modified : `DatasetRef`
A new dataset reference with updated attributes.
"""
if storage_class is None:
datasetType = self.datasetType
else:
datasetType = self.datasetType.overrideStorageClass(storage_class)
if run is None:
run = self.run
# Do not regenerate dataset ID if run is the same.
if id is None:
id = self.id
return DatasetRef(
datasetType=self.datasetType.overrideStorageClass(storageClass),
datasetType=datasetType,
dataId=self.dataId,
id=self.id,
run=self.run,
run=run,
id=id,
conform=False,
)

Expand Down
21 changes: 21 additions & 0 deletions tests/test_datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -643,6 +643,27 @@ def testOverrideStorageClass(self) -> None:
# of "object" which is compatible with everything.
ref_new.overrideStorageClass(incompatible_sc)

def testReplace(self) -> None:
"""Test for `DatasetRef.replace` method."""
ref = DatasetRef(self.datasetType, self.dataId, run="somerun")

ref2 = ref.replace(run="somerun2")
self.assertEqual(ref2.run, "somerun2")
self.assertIsNotNone(ref2.id)
self.assertNotEqual(ref2.id, ref.id)

ref3 = ref.replace(run="somerun3", id=ref2.id)
self.assertEqual(ref3.run, "somerun3")
self.assertEqual(ref3.id, ref2.id)

ref4 = ref.replace(id=ref2.id)
self.assertEqual(ref4.run, "somerun")
self.assertEqual(ref4.id, ref2.id)

ref5 = ref.replace()
self.assertEqual(ref5.run, "somerun")
self.assertEqual(ref5, ref)

def testPickle(self) -> None:
ref = DatasetRef(self.datasetType, self.dataId, run="somerun")
s = pickle.dumps(ref)
Expand Down

0 comments on commit a900423

Please sign in to comment.