lsst-pipe-base 29.2025.3900__py3-none-any.whl → 29.2025.4100__py3-none-any.whl

lsst/pipe/base/tests/mocks/{_in_memory_repo.py → _repo.py}

@@ -27,18 +27,31 @@
 
 from __future__ import annotations
 
-__all__ = ("InMemoryRepo",)
+__all__ = ("DirectButlerRepo", "InMemoryRepo", "MockRepo")
 
-from collections.abc import Iterable, Mapping
+import tempfile
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Iterator, Mapping
+from contextlib import contextmanager
 from typing import Any
 
-from lsst.daf.butler import CollectionType, DataCoordinate, DatasetRef, DatasetType, RegistryConfig
+from lsst.daf.butler import (
+    Butler,
+    CollectionType,
+    DataCoordinate,
+    DatasetRef,
+    DatasetType,
+    DimensionConfig,
+    LimitedButler,
+    RegistryConfig,
+)
 from lsst.daf.butler.tests.utils import create_populated_sqlite_registry
 from lsst.resources import ResourcePath, ResourcePathExpression
 from lsst.sphgeom import RangeSet
 
 from ...all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
 from ...pipeline_graph import PipelineGraph
+from ...quantum_graph import PredictedQuantumGraph
 from ...single_quantum_executor import SingleQuantumExecutor
 from ..in_memory_limited_butler import InMemoryLimitedButler
 from ._pipeline_task import (
@@ -49,16 +62,13 @@ from ._pipeline_task import (
 from ._storage_class import MockDataset, is_mock_name
 
 
-class InMemoryRepo:
-    """A test helper that simulates a butler repository for task execution
-    without any disk I/O.
+class MockRepo(ABC):
+    """A test helper that populates a butler repository for task execution.
 
     Parameters
     ----------
-    *args : `str` or `lsst.resources.ResourcePath`
-        Butler YAML import files to load into the test repository.
-    registry_config : `lsst.daf.butler.RegistryConfig`, optional
-        Registry configuration for the repository.
+    butler : `lsst.daf.butler.Butler`
+        Butler to use for at least quantum graph building.  Must be writeable.
     input_run : `str`, optional
         Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
         used as an input to quantum graph generation.  Input datasets created
@@ -67,49 +77,22 @@ class InMemoryRepo:
         Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
         will be the direct input to quantum graph generation.  This always
         includes ``input_run``.
-    output_run : `str`, optional
-        Name of a `~lsst.daf.butler.CollectionType.RUN` collection for
-        execution outputs.
-    use_import_collections_as_input : `bool` `str`, or \
-            `~collections.abc.Iterable` [ `str`], optional
-        Additional collections from YAML import files to include in
-        ``input_chain``, or `True` to include all such collections (in
-        chain-flattened lexicographical order).
-    data_root : convertible to `lsst.resources.ResourcePath`, optional
-        Root directory to join to each element in ``*args``.  Defaults to
-        the `lsst.daf.butler.tests.registry_data` package.
-
-    Notes
-    -----
-    This helper maintains an `..pipeline_graph.PipelineGraph` and a
-    no-datastore butler backed by an in-memory SQLite database for use in
-    quantum graph generation.
+    input_children : `str` or `~collections.abc.Iterable` [ `str`], optional
+        Additional collections to include in ``input_chain``.
     """
 
     def __init__(
         self,
-        *args: str | ResourcePath,
-        registry_config: RegistryConfig | None = None,
+        butler: Butler,
         input_run: str = "input_run",
         input_chain: str = "input_chain",
-        output_run: str = "output_run",
-        use_import_collections_as_input: bool | str | Iterable[str] = True,
-        data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
+        input_children: Iterable[str] = (),
     ):
-        if data_root is not None:
-            data_root = ResourcePath(data_root, forceDirectory=True)
-            args = tuple(data_root.join(arg) for arg in args)
-        self.butler = create_populated_sqlite_registry(*args, registry_config=registry_config)
+        self.butler = butler
         input_chain_definition = [input_run]
-        if use_import_collections_as_input:
-            if use_import_collections_as_input is True:
-                use_import_collections_as_input = sorted(
-                    self.butler.collections.query("*", flatten_chains=True)
-                )
-            input_chain_definition += list(use_import_collections_as_input)
+        input_chain_definition.extend(input_children)
         self.input_run = input_run
         self.input_chain = input_chain
-        self.output_run = output_run
         self.butler.collections.register(self.input_run)
         self.butler.collections.register(self.input_chain, CollectionType.CHAINED)
         self.butler.collections.redefine_chain(self.input_chain, input_chain_definition)
@@ -206,9 +189,55 @@ class InMemoryRepo:
             label = f"task_auto{self.last_auto_task_index}"
         self.pipeline_graph.add_task(label, task_class=task_class, config=config)
 
+    def make_quantum_graph(
+        self,
+        *,
+        output: str | None = None,
+        output_run: str = "output_run",
+        insert_mocked_inputs: bool = True,
+        register_output_dataset_types: bool = True,
+    ) -> PredictedQuantumGraph:
+        """Make a quantum graph from the pipeline task and internal data
+        repository.
+
+        Parameters
+        ----------
+        output : `str` or `None`, optional
+            Name of the output chained collection to embed within the quantum
+            graph.  Note that this does not actually create this collection.
+        output_run : `str`, optional
+            Name of the `~lsst.daf.butler.CollectionType.RUN` collection for
+            execution outputs.  Note that this does not actually create this
+            collection.
+        insert_mocked_inputs : `bool`, optional
+            Whether to automatically insert datasets for all overall inputs to
+            the pipeline graph whose dataset types have not already been
+            registered.  If set to `False`, inputs must be provided by imported
+            YAML files or explicit calls to `insert_datasets`, which provides
+            more fine-grained control over the data IDs of the datasets.
+        register_output_dataset_types : `bool`, optional
+            If `True`, register all output dataset types.
+
+        Returns
+        -------
+        qg : `..quantum_graph.PredictedQuantumGraph`
+            Quantum graph.  Datastore records will not be attached, since the
+            test helper does not actually have a datastore.
+        """
+        return (
+            self.make_quantum_graph_builder(
+                insert_mocked_inputs=insert_mocked_inputs,
+                register_output_dataset_types=register_output_dataset_types,
+                output_run=output_run,
+            )
+            .finish(output=output, attach_datastore_records=False)
+            .assemble()
+        )
+
     def make_quantum_graph_builder(
         self,
         *,
+        output_run: str = "output_run",
         insert_mocked_inputs: bool = True,
         register_output_dataset_types: bool = True,
     ) -> AllDimensionsQuantumGraphBuilder:
@@ -217,6 +246,10 @@ class InMemoryRepo:
 
         Parameters
         ----------
+        output_run : `str`, optional
+            Name of the `~lsst.daf.butler.CollectionType.RUN` collection for
+            execution outputs.  Note that this does not actually create this
+            collection.
         insert_mocked_inputs : `bool`, optional
             Whether to automatically insert datasets for all overall inputs to
             the pipeline graph whose dataset types have not already been
@@ -244,7 +277,7 @@ class InMemoryRepo:
             self.pipeline_graph,
             self.butler,
             input_collections=[self.input_chain],
-            output_run=self.output_run,
+            output_run=output_run,
         )
         if register_output_dataset_types:
             self.pipeline_graph.register_dataset_types(self.butler)
@@ -311,6 +344,117 @@ class InMemoryRepo:
                 )
         else:
             data_ids = self.butler.query_data_ids(dimensions, *args, **kwargs, explain=False)
+        return self._insert_datasets_impl(dataset_type, data_ids)
+
+    @abstractmethod
+    def _insert_datasets_impl(
+        self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
+    ) -> list[DatasetRef]:
+        """Insert datasets after their data IDs have been generated.
+
+        Parameters
+        ----------
+        dataset_type : `lsst.daf.butler.DatasetType`
+            Type of the datasets.
+        data_ids : `list` [ `lsst.daf.butler.DataCoordinate` ]
+            Data IDs of all datasets.
+
+        Returns
+        -------
+        refs : `list` [ `lsst.daf.butler.DatasetRef` ]
+            References to the new datasets.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def make_single_quantum_executor(
+        self, qg: PredictedQuantumGraph
+    ) -> tuple[SingleQuantumExecutor, LimitedButler]:
+        """Make a single-quantum executor.
+
+        Parameters
+        ----------
+        qg : `..quantum_graph.PredictedQuantumGraph`
+            Graph whose quanta the executor must be capable of executing.
+
+        Returns
+        -------
+        executor : `..single_quantum_executor.SingleQuantumExecutor`
+            An executor for a single quantum.
+        butler : `lsst.daf.butler.LimitedButler`
+            The butler that the executor will write to.
+        """
+        raise NotImplementedError()
+
+
+class InMemoryRepo(MockRepo):
+    """A test helper that simulates a butler repository for task execution
+    without any disk I/O.
+
+    Parameters
+    ----------
+    *args : `str` or `lsst.resources.ResourcePath`
+        Butler YAML import files to load into the test repository.
+    registry_config : `lsst.daf.butler.RegistryConfig`, optional
+        Registry configuration for the repository.
+    dimension_config : `lsst.daf.butler.DimensionConfig`, optional
+        Dimension universe configuration for the repository.
+    input_run : `str`, optional
+        Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
+        used as an input to quantum graph generation.  Input datasets created
+        by the helper are added to this collection.
+    input_chain : `str`, optional
+        Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
+        will be the direct input to quantum graph generation.  This always
+        includes ``input_run``.
+    use_import_collections_as_input : `bool`, `str`, or \
+            `~collections.abc.Iterable` [ `str`], optional
+        Additional collections from YAML import files to include in
+        ``input_chain``, or `True` to include all such collections (in
+        chain-flattened lexicographical order).
+    data_root : convertible to `lsst.resources.ResourcePath`, optional
+        Root directory to join to each element in ``*args``.  Defaults to
+        the `lsst.daf.butler.tests.registry_data` package.
+
+    Notes
+    -----
+    This helper maintains an `..pipeline_graph.PipelineGraph` and a
+    no-datastore butler backed by an in-memory SQLite database for use in
+    quantum graph generation.  It creates a separate in-memory limited butler
+    for execution as needed.
+    """
+
+    def __init__(
+        self,
+        *args: str | ResourcePath,
+        registry_config: RegistryConfig | None = None,
+        dimension_config: DimensionConfig | None = None,
+        input_run: str = "input_run",
+        input_chain: str = "input_chain",
+        use_import_collections_as_input: bool | str | Iterable[str] = True,
+        data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
+    ):
+        if data_root is not None:
+            data_root = ResourcePath(data_root, forceDirectory=True)
+            args = tuple(data_root.join(arg) for arg in args)
+        butler = create_populated_sqlite_registry(
+            *args, registry_config=registry_config, dimension_config=dimension_config
+        )
+        if use_import_collections_as_input:
+            if use_import_collections_as_input is True:
+                use_import_collections_as_input = sorted(butler.collections.query("*", flatten_chains=True))
+        else:
+            use_import_collections_as_input = ()
+        super().__init__(
+            butler,
+            input_run=input_run,
+            input_chain=input_chain,
+            input_children=list(use_import_collections_as_input),
+        )
+
+    def _insert_datasets_impl(
+        self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
+    ) -> list[DatasetRef]:
         return self.butler.registry.insertDatasets(dataset_type, data_ids, run=self.input_run)
 
     def make_limited_butler(self) -> InMemoryLimitedButler:
@@ -343,9 +487,16 @@ class InMemoryRepo:
                 )
         return butler
 
-    def make_single_quantum_executor(self) -> tuple[SingleQuantumExecutor, InMemoryLimitedButler]:
+    def make_single_quantum_executor(
+        self, qg: PredictedQuantumGraph | None = None
+    ) -> tuple[SingleQuantumExecutor, InMemoryLimitedButler]:
         """Make a single-quantum executor backed by a new limited butler.
 
+        Parameters
+        ----------
+        qg : `..quantum_graph.PredictedQuantumGraph`
+            Ignored by this implementation.
+
         Returns
         -------
         executor : `..single_quantum_executor.SingleQuantumExecutor`
@@ -355,3 +506,131 @@ class InMemoryRepo:
         """
         butler = self.make_limited_butler()
         return SingleQuantumExecutor(limited_butler_factory=butler.factory), butler
+
+
+class DirectButlerRepo(MockRepo):
+    """A test helper for task execution backed by a local direct butler.
+
+    Parameters
+    ----------
+    butler : `lsst.daf.butler.direct_butler.DirectButler`
+        Butler to write to.
+    *args : `str` or `lsst.resources.ResourcePath`
+        Butler YAML import files to load into the test repository.
+    input_run : `str`, optional
+        Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
+        used as an input to quantum graph generation.  Input datasets created
+        by the helper are added to this collection.
+    input_chain : `str`, optional
+        Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
+        will be the direct input to quantum graph generation.  This always
+        includes ``input_run``.
+    use_import_collections_as_input : `bool`, `str`, or \
+            `~collections.abc.Iterable` [ `str`], optional
+        Additional collections from YAML import files to include in
+        ``input_chain``, or `True` to include all such collections (in
+        chain-flattened lexicographical order).
+    data_root : convertible to `lsst.resources.ResourcePath`, optional
+        Root directory to join to each element in ``*args``.  Defaults to
+        the `lsst.daf.butler.tests.registry_data` package.
+
+    Notes
+    -----
+    This helper maintains an `..pipeline_graph.PipelineGraph` and a
+    no-datastore butler backed by an in-memory SQLite database for use in
+    quantum graph generation.  It creates a separate in-memory limited butler
+    for execution as needed.
+    """
+
+    def __init__(
+        self,
+        butler: Butler,
+        *args: str | ResourcePath,
+        input_run: str = "input_run",
+        input_chain: str = "input_chain",
+        use_import_collections_as_input: bool | str | Iterable[str] = True,
+        data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
+    ):
+        if data_root is not None:
+            data_root = ResourcePath(data_root, forceDirectory=True)
+            args = tuple(data_root.join(arg) for arg in args)
+        for arg in args:
+            butler.import_(filename=arg)
+        if use_import_collections_as_input:
+            if use_import_collections_as_input is True:
+                use_import_collections_as_input = sorted(butler.collections.query("*", flatten_chains=True))
+        else:
+            use_import_collections_as_input = ()
+        super().__init__(
+            butler,
+            input_run=input_run,
+            input_chain=input_chain,
+            input_children=list(use_import_collections_as_input),
+        )
+
+    @classmethod
+    @contextmanager
+    def make_temporary(
+        cls,
+        *args: str | ResourcePath,
+        input_run: str = "input_run",
+        input_chain: str = "input_chain",
+        use_import_collections_as_input: bool | str | Iterable[str] = True,
+        data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
+        **kwargs: Any,
+    ) -> Iterator[tuple[DirectButlerRepo, str]]:
+        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as root:
+            config = Butler.makeRepo(root, **kwargs)
+            butler = Butler.from_config(config, writeable=True)
+            yield (
+                cls(
+                    butler,
+                    *args,
+                    input_run=input_run,
+                    input_chain=input_chain,
+                    use_import_collections_as_input=use_import_collections_as_input,
+                    data_root=data_root,
+                ),
+                root,
+            )
+
+    def _insert_datasets_impl(
+        self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
+    ) -> list[DatasetRef]:
+        if is_mock_name(dataset_type.storageClass_name):
+            refs: list[DatasetRef] = []
+            for data_id in data_ids:
+                data_id = self.butler.registry.expandDataId(data_id)
+                ref = DatasetRef(dataset_type, data_id, run=self.input_run)
+                self.butler.put(
+                    MockDataset(
+                        dataset_id=ref.id,
+                        dataset_type=ref.datasetType.to_simple(),
+                        data_id=dict(ref.dataId.mapping),
+                        run=ref.run,
+                    ),
+                    ref,
+                )
+                refs.append(ref)
+            return refs
+        else:
+            return self.butler.registry.insertDatasets(dataset_type, data_ids, run=self.input_run)
+
+    def make_single_quantum_executor(
+        self, qg: PredictedQuantumGraph | None = None
+    ) -> tuple[SingleQuantumExecutor, Butler]:
+        """Make a single-quantum executor backed by a new limited butler.
+
+        Parameters
+        ----------
+        qg : `..quantum_graph.PredictedQuantumGraph`
+            Ignored by this implementation.
+
+        Returns
+        -------
+        executor : `..single_quantum_executor.SingleQuantumExecutor`
+            An executor for a single quantum.
+        butler : `lsst.daf.butler.Butler`
+            The butler that the executor will write to.
+        """
+        return SingleQuantumExecutor(limited_butler_factory=None, butler=self.butler), self.butler

lsst/pipe/base/tests/mocks/_storage_class.py

@@ -173,6 +173,12 @@ class MockDataset(pydantic.BaseModel):
     parameters: dict[str, str] | None = None
     """`repr` of all parameters applied when reading this dataset."""
 
+    int_value: int | None = None
+    """An arbitrary integer value stored in the mock dataset."""
+
+    str_value: int | None = None
+    """An arbitrary string value stored in the mock dataset."""
+
     @property
     def storage_class(self) -> str:
         return cast(str, self.dataset_type.storageClass)
@@ -236,6 +242,21 @@ class MockDataset(pydantic.BaseModel):
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
 
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)
+
 
 class ConvertedUnmockedDataset(pydantic.BaseModel):
     """A marker class that represents a conversion from a regular in-memory
@@ -275,6 +296,21 @@ class ConvertedUnmockedDataset(pydantic.BaseModel):
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
 
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)
+
 
 class MockDatasetQuantum(pydantic.BaseModel):
     """Description of the quantum that produced a mock dataset.
@@ -325,6 +361,21 @@ class MockDatasetQuantum(pydantic.BaseModel):
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
 
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)
+
 
 MockDataset.model_rebuild()
 

lsst/pipe/base/tests/simpleQGraph.py

@@ -49,7 +49,13 @@ from .. import connectionTypes as cT
 from .._instrument import Instrument
 from ..all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
 from ..all_dimensions_quantum_graph_builder import DatasetQueryConstraintVariant as DSQVariant
-from ..automatic_connection_constants import PACKAGES_INIT_OUTPUT_NAME, PACKAGES_INIT_OUTPUT_STORAGE_CLASS
+from ..automatic_connection_constants import (
+    CONFIG_INIT_OUTPUT_CONNECTION_NAME,
+    LOG_OUTPUT_CONNECTION_NAME,
+    METADATA_OUTPUT_CONNECTION_NAME,
+    PACKAGES_INIT_OUTPUT_NAME,
+    PACKAGES_INIT_OUTPUT_STORAGE_CLASS,
+)
 from ..config import PipelineTaskConfig
 from ..connections import PipelineTaskConnections
 from ..graph import QuantumGraph
@@ -361,13 +367,13 @@ def populateButler(
         if run is not None:
             butler.registry.registerRun(run)
         for dsType in dsTypes:
-            if dsType == "packages":
+            if dsType == PACKAGES_INIT_OUTPUT_NAME:
                 # Version is intentionally inconsistent.
                 # Dict is convertible to Packages if Packages is installed.
                 data: Any = {"python": "9.9.99"}
                 butler.put(data, dsType, run=run)
             else:
-                if dsType.endswith("_config"):
+                if dsType.endswith(CONFIG_INIT_OUTPUT_CONNECTION_NAME):
                     # find a config from matching task name or make a new one
                     taskLabel, _, _ = dsType.rpartition("_")
                     task_node = pipeline_graph.tasks.get(taskLabel)
@@ -375,9 +381,9 @@ def populateButler(
                         data = task_node.config
                     else:
                         data = AddTaskConfig()
-                elif dsType.endswith("_metadata"):
+                elif dsType.endswith(METADATA_OUTPUT_CONNECTION_NAME):
                     data = _TASK_FULL_METADATA_TYPE()
-                elif dsType.endswith("_log"):
+                elif dsType.endswith(LOG_OUTPUT_CONNECTION_NAME):
                     data = ButlerLogRecords.from_records([])
                 else:
                     data = numpy.array([0.0, 1.0, 2.0, 5.0])

lsst/pipe/base/version.py

@@ -1,2 +1,2 @@
 __all__ = ["__version__"]
-__version__ = "29.2025.3900"
+__version__ = "29.2025.4100"

{lsst_pipe_base-29.2025.3900.dist-info → lsst_pipe_base-29.2025.4100.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lsst-pipe-base
-Version: 29.2025.3900
+Version: 29.2025.4100
 Summary: Pipeline infrastructure for the Rubin Science Pipelines.
 Author-email: Rubin Observatory Data Management <dm-admin@lists.lsst.org>
 License: BSD 3-Clause License
@@ -31,6 +31,7 @@ Requires-Dist: wcwidth
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: numpy>=1.17
 Requires-Dist: frozendict
+Requires-Dist: zstandard<0.24,>=0.23.0
 Provides-Extra: test
 Requires-Dist: pytest>=3.2; extra == "test"
 Provides-Extra: mermaid