lsst-daf-butler 29.0.1rc1__py3-none-any.whl → 29.1.0__py3-none-any.whl

lsst/daf/butler/_quantum_backed.py

@@ -43,6 +43,7 @@ import pydantic
 from lsst.resources import ResourcePath, ResourcePathExpression
 
 from ._butler_config import ButlerConfig
+from ._butler_metrics import ButlerMetrics
 from ._config import Config
 from ._dataset_provenance import DatasetProvenance
 from ._dataset_ref import DatasetId, DatasetRef
@@ -118,6 +119,8 @@ class QuantumBackedButler(LimitedButler):
         Object managing all storage class definitions.
     dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`]
         The registry dataset type definitions, indexed by name.
+    metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
+        Metrics object for tracking butler statistics.
 
     Notes
     -----
@@ -164,6 +167,7 @@ class QuantumBackedButler(LimitedButler):
         datastore: Datastore,
         storageClasses: StorageClassFactory,
         dataset_types: Mapping[str, DatasetType] | None = None,
+        metrics: ButlerMetrics | None = None,
     ):
         self._dimensions = dimensions
         self._predicted_inputs = set(predicted_inputs)
@@ -175,6 +179,7 @@ class QuantumBackedButler(LimitedButler):
         self._datastore = datastore
         self.storageClasses = storageClasses
         self._dataset_types: Mapping[str, DatasetType] = {}
+        self._metrics = metrics if metrics is not None else ButlerMetrics()
         if dataset_types is not None:
             self._dataset_types = dataset_types
         self._datastore.set_retrieve_dataset_type_method(self._retrieve_dataset_type)
@@ -190,6 +195,7 @@ class QuantumBackedButler(LimitedButler):
         BridgeManagerClass: type[DatastoreRegistryBridgeManager] = MonolithicDatastoreRegistryBridgeManager,
         search_paths: list[str] | None = None,
         dataset_types: Mapping[str, DatasetType] | None = None,
+        metrics: ButlerMetrics | None = None,
     ) -> QuantumBackedButler:
         """Construct a new `QuantumBackedButler` from repository configuration
         and helper types.
@@ -219,6 +225,8 @@ class QuantumBackedButler(LimitedButler):
         dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`], \
                 optional
             Mapping of the dataset type name to its registry definition.
+        metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
+            Metrics object for gathering butler statistics.
         """
         predicted_inputs = [ref.id for ref in itertools.chain.from_iterable(quantum.inputs.values())]
         predicted_inputs += [ref.id for ref in quantum.initInputs.values()]
@@ -234,6 +242,7 @@ class QuantumBackedButler(LimitedButler):
             BridgeManagerClass=BridgeManagerClass,
             search_paths=search_paths,
             dataset_types=dataset_types,
+            metrics=metrics,
         )
 
     @classmethod
@@ -249,6 +258,7 @@ class QuantumBackedButler(LimitedButler):
         BridgeManagerClass: type[DatastoreRegistryBridgeManager] = MonolithicDatastoreRegistryBridgeManager,
         search_paths: list[str] | None = None,
         dataset_types: Mapping[str, DatasetType] | None = None,
+        metrics: ButlerMetrics | None = None,
     ) -> QuantumBackedButler:
         """Construct a new `QuantumBackedButler` from sets of input and output
         dataset IDs.
@@ -281,6 +291,8 @@ class QuantumBackedButler(LimitedButler):
         dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`], \
                 optional
             Mapping of the dataset type name to its registry definition.
+        metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
+            Metrics object for gathering butler statistics.
         """
         return cls._initialize(
             config=config,
@@ -293,6 +305,7 @@ class QuantumBackedButler(LimitedButler):
             BridgeManagerClass=BridgeManagerClass,
             search_paths=search_paths,
             dataset_types=dataset_types,
+            metrics=metrics,
         )
 
     @classmethod
@@ -309,6 +322,7 @@ class QuantumBackedButler(LimitedButler):
         BridgeManagerClass: type[DatastoreRegistryBridgeManager] = MonolithicDatastoreRegistryBridgeManager,
         search_paths: list[str] | None = None,
         dataset_types: Mapping[str, DatasetType] | None = None,
+        metrics: ButlerMetrics | None = None,
     ) -> QuantumBackedButler:
         """Initialize quantum-backed butler.
 
@@ -341,6 +355,8 @@ class QuantumBackedButler(LimitedButler):
             Additional search paths for butler configuration.
         dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`]
             Mapping of the dataset type name to its registry definition.
+        metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
+            Metrics object for gathering butler statistics.
         """
         butler_config = ButlerConfig(config, searchPaths=search_paths)
         butler_root = butler_config.get("root", butler_config.configDir)
@@ -373,6 +389,7 @@ class QuantumBackedButler(LimitedButler):
             datastore,
             storageClasses=storageClasses,
             dataset_types=dataset_types,
+            metrics=metrics,
         )
 
     def _retrieve_dataset_type(self, name: str) -> DatasetType | None:
@@ -459,8 +476,9 @@ class QuantumBackedButler(LimitedButler):
         # Docstring inherited.
         if ref.id not in self._predicted_outputs:
             raise RuntimeError("Cannot `put` dataset that was not predicted as an output.")
-        self._datastore.put(obj, ref, provenance=provenance)
-        self._actual_output_refs.add(ref)
+        with self._metrics.instrument_put(log=_LOG, msg="Put QBB dataset"):
+            self._datastore.put(obj, ref, provenance=provenance)
+            self._actual_output_refs.add(ref)
         return ref
 
     def pruneDatasets(
@@ -498,8 +516,9 @@ class QuantumBackedButler(LimitedButler):
                 self._actual_output_refs.discard(ref)
 
         if unstore:
-            # Point of no return for removing artifacts
-            self._datastore.emptyTrash()
+            # Point of no return for removing artifacts. Only try to remove
+            # refs associated with this pruning.
+            self._datastore.emptyTrash(refs=refs)
 
     def retrieve_artifacts_zip(
         self,

lsst/daf/butler/arrow_utils.py

@@ -383,7 +383,7 @@ class _ToArrowTimespan(ToArrow):
         # Docstring inherited.
         return TimespanArrowType()
 
-    def append(self, value: Timespan | None, column: list[pa.StructScalar | None]) -> None:
+    def append(self, value: Timespan | None, column: list[dict[str, int] | None]) -> None:
         # Docstring inherited.
         column.append({"begin_nsec": value.nsec[0], "end_nsec": value.nsec[1]} if value is not None else None)
 
@@ -432,12 +432,10 @@ class _ToArrowDateTime(ToArrow):
 
 @final
 class UUIDArrowType(pa.ExtensionType):
-    """An Arrow extension type for `astropy.time.Time`, stored as TAI
-    nanoseconds since 1970-01-01.
-    """
+    """An Arrow extension type for `uuid.UUID`, stored as 16 bytes."""
 
     def __init__(self) -> None:
-        super().__init__(_ToArrowTimespan.storage_type, "astropy.time.Time")
+        super().__init__(_ToArrowUUID.storage_type, "uuid.UUID")
 
     def __arrow_ext_serialize__(self) -> bytes:
         return b""
@@ -458,7 +456,7 @@ class UUIDArrowScalar(pa.ExtensionScalar):
     instance.
     """
 
-    def as_py(self) -> astropy.time.Time:
+    def as_py(self, **_unused: Any) -> uuid.UUID:
         return uuid.UUID(bytes=self.value.as_py())
 
 
@@ -487,7 +485,7 @@ class RegionArrowScalar(pa.ExtensionScalar):
     Use the standard `as_py` method to convert to an actual region.
     """
 
-    def as_py(self) -> Region:
+    def as_py(self, **_unused: Any) -> Region:
         return Region.decode(self.value.as_py())
 
 
@@ -516,7 +514,7 @@ class TimespanArrowScalar(pa.ExtensionScalar):
     Use the standard `as_py` method to convert to an actual timespan.
     """
 
-    def as_py(self) -> Timespan | None:
+    def as_py(self, **_unused: Any) -> Timespan | None:
         if self.value is None:
             return None
         else:
@@ -554,7 +552,7 @@ class DateTimeArrowScalar(pa.ExtensionScalar):
     instance.
     """
 
-    def as_py(self) -> astropy.time.Time:
+    def as_py(self, **_unused: Any) -> astropy.time.Time:
         return TimeConverter().nsec_to_astropy(self.value.as_py())
 
 

lsst/daf/butler/cli/butler.py

@@ -102,7 +102,7 @@ class PluginCommand:
     """Where the command came from (`str`)."""
 
 
-class LoaderCLI(click.MultiCommand, abc.ABC):
+class LoaderCLI(click.Group, abc.ABC):
     """Extends `click.MultiCommand`, which dispatches to subcommands, to load
     subcommands at runtime.
 

lsst/daf/butler/cli/cmd/_remove_runs.py

@@ -81,7 +81,9 @@ def _print_remove(will: bool, runs: Sequence[script.RemoveRun], datasets: Mappin
         else:
             print(run.name)
     print("\n" + willRemoveDatasetsMsg if will else didRemoveDatasetsMsg)
+    total = sum(datasets.values())
     print(", ".join([f"{i[0]}({i[1]})" for i in datasets.items()]))
+    print("Total number of datasets to remove: ", total)
 
 
 def _print_requires_confirmation(runs: Sequence[script.RemoveRun], datasets: Mapping[str, int]) -> None:

lsst/daf/butler/cli/cmd/commands.py

@@ -424,6 +424,21 @@ def prune_datasets(**kwargs: Any) -> None:
         case_sensitive=False,
     ),
 )
+@click.option(
+    "-t",
+    "--show-dataset-types",
+    is_flag=True,
+    help="Also show the dataset types registered within each collection.",
+)
+@click.option(
+    "--exclude-dataset-types",
+    type=click.STRING,
+    multiple=True,
+    default=["*_config,*_log,*_metadata,packages"],
+    callback=split_commas,
+    show_default=True,
+    help="Dataset types (comma-separated) to exclude. Only valid with --show-dataset-types.",
+)
 @options_file_option()
 def query_collections(*args: Any, **kwargs: Any) -> None:
     """Get the collections whose names match an expression."""
@@ -454,7 +469,7 @@ def query_dataset_types(*args: Any, **kwargs: Any) -> None:
     """Get the dataset types in a repository."""
     table = script.queryDatasetTypes(*args, **kwargs)
     if table:
-        table.pprint_all()
+        table.pprint_all(align="<")
     else:
         print("No results. Try --help for more information.")
 
@@ -638,6 +653,9 @@ def retrieve_artifacts(**kwargs: Any) -> None:
 @transfer_option()
 @register_dataset_types_option()
 @transfer_dimensions_option()
+@click.option(
+    "--dry-run/--no-dry-run", default=False, help="Enable dry run mode and do not transfer any datasets."
+)
 @options_file_option()
 def transfer_datasets(**kwargs: Any) -> None:
     """Transfer datasets from a source butler to a destination butler.
@@ -826,6 +844,12 @@ def export_calibs(*args: Any, **kwargs: Any) -> None:
 @repo_argument(required=True)
 @click.argument("zip", required=True)
 @transfer_option()
+@transfer_dimensions_option(
+    default=False, help="Attempt to register missing dimension records during ingest."
+)
+@click.option(
+    "--dry-run/--no-dry-run", default=False, help="Enable dry run mode and do not ingest any datasets."
+)
 def ingest_zip(**kwargs: Any) -> None:
     """Ingest a Zip file created by retrieve-artifacts.
 

lsst/daf/butler/cli/utils.py

@@ -55,6 +55,7 @@ __all__ = (
 )
 
 
+import importlib.metadata
 import itertools
 import logging
 import os
@@ -76,6 +77,7 @@ import click
 import click.exceptions
 import click.testing
 import yaml
+from packaging.version import Version
 
 from lsst.utils.iteration import ensure_iterable
 
@@ -87,6 +89,12 @@ if TYPE_CHECKING:
 
     from lsst.daf.butler import Dimension
 
+_click_version = Version(importlib.metadata.version("click"))
+if _click_version >= Version("8.2.0"):
+    _click_make_metavar_has_context = True
+else:
+    _click_make_metavar_has_context = False
+
 log = logging.getLogger(__name__)
 
 # This is used as the metavar argument to Options that accept multiple string
@@ -741,9 +749,16 @@ class MWPath(click.Path):
 class MWOption(click.Option):
     """Overrides click.Option with desired behaviors."""
 
-    def make_metavar(self) -> str:
+    def make_metavar(self, ctx: click.Context | None = None) -> str:
         """Make the metavar for the help menu.
 
+        Parameters
+        ----------
+        ctx : `click.Context` or `None`
+            Context from the command.
+
+        Notes
+        -----
         Overrides `click.Option.make_metavar`.
         Adds a space and an ellipsis after the metavar name if
         the option accepts multiple inputs, otherwise defers to the base
@@ -758,7 +773,10 @@ class MWOption(click.Option):
         transformation that must apply to all types should be applied in
         get_help_record.
         """
-        metavar = super().make_metavar()
+        if _click_make_metavar_has_context:
+            metavar = super().make_metavar(ctx=ctx)  # type: ignore
+        else:
+            metavar = super().make_metavar()  # type: ignore
         if self.multiple and self.nargs == 1:
             metavar += " ..."
         elif self.nargs != 1:
@@ -769,9 +787,16 @@ class MWOption(click.Option):
 class MWArgument(click.Argument):
     """Overrides click.Argument with desired behaviors."""
 
-    def make_metavar(self) -> str:
+    def make_metavar(self, ctx: click.Context | None = None) -> str:
         """Make the metavar for the help menu.
 
+        Parameters
+        ----------
+        ctx : `click.Context` or `None`
+            Context from the command.
+
+        Notes
+        -----
         Overrides `click.Option.make_metavar`.
         Always adds a space and an ellipsis (' ...') after the
         metavar name if the option accepts multiple inputs.
@@ -784,7 +809,10 @@ class MWArgument(click.Argument):
         metavar : `str`
             The metavar value.
         """
-        metavar = super().make_metavar()
+        if _click_make_metavar_has_context:
+            metavar = super().make_metavar(ctx=ctx)  # type: ignore
+        else:
+            metavar = super().make_metavar()  # type: ignore
         if self.nargs != 1:
             metavar = f"{metavar[:-3]} ..."
         return metavar
@@ -1029,7 +1057,7 @@ class MWCommand(click.Command):
         return ret
 
     @epilog.setter
-    def epilog(self, val: str) -> None:
+    def epilog(self, val: str | None) -> None:
         self._epilog = val
 
 

lsst/daf/butler/column_spec.py

@@ -39,12 +39,24 @@ __all__ = (
     "StringColumnSpec",
     "TimespanColumnSpec",
     "UUIDColumnSpec",
+    "make_tuple_type_adapter",
 )
 
 import textwrap
 import uuid
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, TypeAlias, Union, final
+from collections.abc import Iterable
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    ClassVar,
+    Literal,
+    Optional,
+    TypeAlias,
+    Union,
+    final,
+)
 
 import astropy.time
 import pyarrow as pa
@@ -54,7 +66,7 @@ from lsst.sphgeom import Region
 
 from . import arrow_utils, ddl
 from ._timespan import Timespan
-from .pydantic_utils import SerializableRegion, SerializableTime
+from .pydantic_utils import SerializableBytesHex, SerializableRegion, SerializableTime
 
 if TYPE_CHECKING:
     from .name_shrinker import NameShrinker
@@ -125,18 +137,6 @@ class ColumnValueSerializer(ABC):
         raise NotImplementedError
 
 
-class _DefaultColumnValueSerializer(ColumnValueSerializer):
-    """Default implementation of serializer for basic types."""
-
-    def serialize(self, value: Any) -> Any:
-        # Docstring inherited.
-        return value
-
-    def deserialize(self, value: Any) -> Any:
-        # Docstring inherited.
-        return value
-
-
 class _TypeAdapterColumnValueSerializer(ColumnValueSerializer):
     """Implementation of serializer that uses pydantic type adapter."""
 
@@ -156,6 +156,8 @@ class _TypeAdapterColumnValueSerializer(ColumnValueSerializer):
 class _BaseColumnSpec(pydantic.BaseModel, ABC):
     """Base class for descriptions of table columns."""
 
+    pytype: ClassVar[type]
+
     name: str = pydantic.Field(description="""Name of the column.""")
 
     doc: str = pydantic.Field(default="", description="Documentation for the column.")
@@ -200,7 +202,6 @@ class _BaseColumnSpec(pydantic.BaseModel, ABC):
         """
         raise NotImplementedError()
 
-    @abstractmethod
     def serializer(self) -> ColumnValueSerializer:
         """Return object that converts values of this column to or from
         serializable format.
@@ -210,7 +211,7 @@ class _BaseColumnSpec(pydantic.BaseModel, ABC):
         serializer : `ColumnValueSerializer`
             A converter instance.
         """
-        raise NotImplementedError()
+        return _TypeAdapterColumnValueSerializer(pydantic.TypeAdapter(self.annotated_type))
 
     def display(self, level: int = 0, tab: str = "  ") -> list[str]:
         """Return a human-reader-focused string description of this column as
@@ -243,6 +244,48 @@ class _BaseColumnSpec(pydantic.BaseModel, ABC):
     def __str__(self) -> str:
         return "\n".join(self.display())
 
+    @property
+    def annotated_type(self) -> Any:
+        """Return a Pydantic-friendly type annotation for this column type.
+
+        Since this is a runtime object and most type annotations must be
+        static, this is really only useful for `pydantic.TypeAdapter`
+        construction and dynamic `pydantic.create_model` construction.
+        """
+        base = self._get_base_annotated_type()
+        if self.nullable:
+            return Optional[base]
+        return base
+
+    @abstractmethod
+    def _get_base_annotated_type(self) -> Any:
+        """Return the base annotated type (not taking into account `nullable`)
+        for this column type.
+        """
+        raise NotImplementedError()
+
+
+def make_tuple_type_adapter(
+    columns: Iterable[ColumnSpec],
+) -> pydantic.TypeAdapter[tuple[Any, ...]]:
+    """Return a `pydantic.TypeAdapter` for a `tuple` with types defined by an
+    iterable of `ColumnSpec` objects.
+
+    Parameters
+    ----------
+    columns : `~collections.abc.Iterable` [ `ColumnSpec` ]
+        Iterable of column specifications.
+
+    Returns
+    -------
+    adapter : `pydantic.TypeAdapter`
+        A Pydantic type adapter for the `tuple` representation of a row with
+        the given columns.
+    """
+    # Static type-checkers don't like this runtime use of static-typing
+    # constructs, but that's how Pydantic works.
+    return pydantic.TypeAdapter(tuple[*[spec.annotated_type for spec in columns]])  # type: ignore
+
 
 @final
 class IntColumnSpec(_BaseColumnSpec):
@@ -256,9 +299,9 @@ class IntColumnSpec(_BaseColumnSpec):
         # Docstring inherited.
         return arrow_utils.ToArrow.for_primitive(self.name, pa.uint64(), nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _DefaultColumnValueSerializer()
+        return pydantic.StrictInt
 
 
 @final
@@ -280,9 +323,9 @@ class StringColumnSpec(_BaseColumnSpec):
         # Docstring inherited.
         return arrow_utils.ToArrow.for_primitive(self.name, pa.string(), nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _DefaultColumnValueSerializer()
+        return pydantic.StrictStr
 
 
 @final
@@ -310,9 +353,9 @@ class HashColumnSpec(_BaseColumnSpec):
             nullable=self.nullable,
         )
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _DefaultColumnValueSerializer()
+        return SerializableBytesHex
 
 
 @final
@@ -328,9 +371,9 @@ class FloatColumnSpec(_BaseColumnSpec):
         assert self.nullable is not None, "nullable=None should be resolved by validators"
         return arrow_utils.ToArrow.for_primitive(self.name, pa.float64(), nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _DefaultColumnValueSerializer()
+        return pydantic.StrictFloat
 
 
 @final
@@ -345,9 +388,9 @@ class BoolColumnSpec(_BaseColumnSpec):
         # Docstring inherited.
         return arrow_utils.ToArrow.for_primitive(self.name, pa.bool_(), nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _DefaultColumnValueSerializer()
+        return pydantic.StrictBool
 
 
 @final
@@ -363,9 +406,9 @@ class UUIDColumnSpec(_BaseColumnSpec):
         assert self.nullable is not None, "nullable=None should be resolved by validators"
         return arrow_utils.ToArrow.for_uuid(self.name, nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _TypeAdapterColumnValueSerializer(pydantic.TypeAdapter(self.pytype))
+        return uuid.UUID
 
 
 @final
@@ -386,9 +429,9 @@ class RegionColumnSpec(_BaseColumnSpec):
         assert self.nullable is not None, "nullable=None should be resolved by validators"
         return arrow_utils.ToArrow.for_region(self.name, nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _TypeAdapterColumnValueSerializer(pydantic.TypeAdapter(SerializableRegion))
+        return SerializableRegion
 
 
 @final
@@ -405,9 +448,9 @@ class TimespanColumnSpec(_BaseColumnSpec):
         # Docstring inherited.
         return arrow_utils.ToArrow.for_timespan(self.name, nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _TypeAdapterColumnValueSerializer(pydantic.TypeAdapter(self.pytype))
+        return Timespan
 
 
 @final
@@ -425,9 +468,9 @@ class DateTimeColumnSpec(_BaseColumnSpec):
         assert self.nullable is not None, "nullable=None should be resolved by validators"
         return arrow_utils.ToArrow.for_datetime(self.name, nullable=self.nullable)
 
-    def serializer(self) -> ColumnValueSerializer:
+    def _get_base_annotated_type(self) -> Any:
         # Docstring inherited.
-        return _TypeAdapterColumnValueSerializer(pydantic.TypeAdapter(SerializableTime))
+        return SerializableTime
 
 
 ColumnSpec = Annotated[

lsst/daf/butler/configs/datastores/formatters.yaml

@@ -94,3 +94,4 @@ Timespan: lsst.daf.butler.formatters.json.JsonFormatter
 RegionTimeInfo: lsst.daf.butler.formatters.json.JsonFormatter
 QPEnsemble: lsst.meas.pz.qp_formatter.QPFormatter
 PZModel: lsst.meas.pz.model_formatter.ModelFormatter
+VisitBackgroundModel: lsst.daf.butler.formatters.json.JsonFormatter

lsst/daf/butler/configs/storageClasses.yaml

@@ -423,3 +423,5 @@ storageClasses:
     pytype: qp.Ensemble
   PZModel:
     pytype: rail.core.model.Model
+  VisitBackgroundModel:
+    pytype: lsst.drp.tasks.fit_visit_background.VisitBackgroundModel

lsst/daf/butler/datastore/__init__.py

@@ -26,3 +26,4 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from ._datastore import *
+from ._transfer import *