diff --git a/doc/changes/DM-47328.api.rst b/doc/changes/DM-47328.api.rst new file mode 100644 index 0000000000..8dcca0b133 --- /dev/null +++ b/doc/changes/DM-47328.api.rst @@ -0,0 +1 @@ +Added ``QuantumBackedButler.retrieve_artifacts`` method to allow dataset artifacts to be retrieved from a graph. diff --git a/python/lsst/daf/butler/_quantum_backed.py b/python/lsst/daf/butler/_quantum_backed.py index 067fb93c45..0f3909dde0 100644 --- a/python/lsst/daf/butler/_quantum_backed.py +++ b/python/lsst/daf/butler/_quantum_backed.py @@ -531,6 +531,53 @@ def retrieve_artifacts_zip( """ return retrieve_and_zip(refs, destination, self._datastore.retrieveArtifacts, overwrite) + def retrieve_artifacts( + self, + refs: Iterable[DatasetRef], + destination: ResourcePathExpression, + transfer: str = "auto", + preserve_path: bool = True, + overwrite: bool = False, + ) -> list[ResourcePath]: + """Retrieve the artifacts associated with the supplied refs. + + Parameters + ---------- + refs : iterable of `DatasetRef` + The datasets for which artifacts are to be retrieved. + A single ref can result in multiple artifacts. The refs must + be resolved. + destination : `lsst.resources.ResourcePath` or `str` + Location to write the artifacts. + transfer : `str`, optional + Method to use to transfer the artifacts. Must be one of the options + supported by `~lsst.resources.ResourcePath.transfer_from()`. + "move" is not allowed. + preserve_path : `bool`, optional + If `True` the full path of the artifact within the datastore + is preserved. If `False` the final file component of the path + is used. + overwrite : `bool`, optional + If `True` allow transfers to overwrite existing files at the + destination. + + Returns + ------- + targets : `list` of `lsst.resources.ResourcePath` + URIs of file artifacts in destination location. Order is not + preserved. + """ + outdir = ResourcePath(destination) + artifact_map = self._datastore.retrieveArtifacts( + refs, + outdir, + transfer=transfer, + preserve_path=preserve_path, + overwrite=overwrite, + write_index=True, + ) + return list(artifact_map) + def extract_provenance_data(self) -> QuantumProvenanceData: """Extract provenance information and datastore records from this butler. diff --git a/tests/test_quantumBackedButler.py b/tests/test_quantumBackedButler.py index 97e5276057..97c0360adb 100644 --- a/tests/test_quantumBackedButler.py +++ b/tests/test_quantumBackedButler.py @@ -230,6 +230,15 @@ def test_getput(self) -> None: self.assertEqual(set(zip_refs), set(self.output_refs)) self.assertEqual(len(index.artifact_map), 4) # Count number of artifacts in Zip. + # Retrieve them to a directory. + with tempfile.TemporaryDirectory() as tmpdir: + retrieved = qbb.retrieve_artifacts(self.output_refs, destination=tmpdir, preserve_path=False) + self.assertEqual(len(retrieved), 4) + self.assertTrue(retrieved[0].exists()) + with open(os.path.join(tmpdir, ZipIndex.index_name)) as zf: + index = ZipIndex.model_validate_json(zf.read()) + self.assertEqual(len(index.artifact_map), len(retrieved)) + def test_getDeferred(self) -> None: """Test for getDeferred method""" quantum = self.make_quantum()