From cda19feb080f8f0746295532a9a34745673ad6b6 Mon Sep 17 00:00:00 2001
From: Andy Salnikov <salnikov@slac.stanford.edu>
Date: Thu, 17 Aug 2023 15:56:04 -0700
Subject: [PATCH] Add methods to update run collection name in a quantum
 (DM-40392)

QuantumGraph.updateRun() method needs a better implementation and
these new methods in DatasetRef and Quantum classes are for supporting
that implementation.
---
 python/lsst/daf/butler/core/datasets/ref.py | 43 +++++++++++++++++++--
 tests/test_datasets.py                      | 21 ++++++++++
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/python/lsst/daf/butler/core/datasets/ref.py b/python/lsst/daf/butler/core/datasets/ref.py
index cd20fca164..38b8d9ef3f 100644
--- a/python/lsst/daf/butler/core/datasets/ref.py
+++ b/python/lsst/daf/butler/core/datasets/ref.py
@@ -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,
         )
 
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index ac8196e4d8..07a2f0af9a 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -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)