Add method to replace DatasetRef data ID with expanded one.

We need this to fix previous loophole usage in queryDatasets.
lsst · Dec 5, 2019 · c48e0ea · c48e0ea
1 parent 4619ccd
commit c48e0ea
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/python/lsst/daf/butler/core/datasets/ref.py b/python/lsst/daf/butler/core/datasets/ref.py
@@ -26,7 +26,7 @@
 from typing import Any, Dict, Mapping, Optional, Tuple
 
 from types import MappingProxyType
-from ..dimensions import DataCoordinate, DimensionGraph
+from ..dimensions import DataCoordinate, DimensionGraph, ExpandedDataCoordinate
 from ..run import Run
 from ..configSupport import LookupKey
 from ..utils import immutable
@@ -183,6 +183,25 @@ def unresolved(self) -> DatasetRef:
         """
         return DatasetRef(datasetType=self.datasetType, dataId=self.dataId, hash=self.hash, conform=False)
 
+    def expanded(self, dataId: ExpandedDataCoordinate) -> DatasetRef:
+        """Return a new `DatasetRef` with the given expanded data ID.
+
+        Parameters
+        ----------
+        dataId : `ExpandedDataCoordinate`
+            Data ID for the new `DatasetRef`.  Must compare equal to the
+            original data ID.
+
+        Returns
+        -------
+        ref : `DatasetRef`
+            A new `DatasetRef` with the given data ID.
+        """
+        assert dataId == self.dataId
+        return DatasetRef(datasetType=self.datasetType, dataId=dataId,
+                          id=self.id, run=self.run, hash=self.hash, components=self.components,
+                          conform=False)
+
     def isComponent(self) -> bool:
         """Boolean indicating whether this `DatasetRef` refers to a
         component of a composite.

diff --git a/python/lsst/daf/butler/registries/sqlRegistry.py b/python/lsst/daf/butler/registries/sqlRegistry.py
@@ -928,8 +928,7 @@ def queryDatasets(self, datasetType: DatasetTypeExpression, *,
             if expand:
                 for ref in bestRefs.values():
                     dataId = self.expandDataId(ref.dataId, records=standardizedDataId.records)
-                    ref._dataId = dataId  # TODO: add semi-public API for this?
-                    yield ref
+                    yield ref.expanded(dataId)
             else:
                 yield from bestRefs.values()