lsst-pipe-base 30.0.0rc3__py3-none-any.whl → 30.0.1__py3-none-any.whl

lsst/pipe/base/quantum_graph/_common.py

@@ -50,9 +50,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Self,
-    TypeAlias,
     TypedDict,
-    TypeVar,
 )
 
 import networkx
@@ -81,18 +79,16 @@ if TYPE_CHECKING:
 # These aliases make it a lot easier how the various pydantic models are
 # structured, but they're too verbose to be worth exporting to code outside the
 # quantum_graph subpackage.
-TaskLabel: TypeAlias = str
-DatasetTypeName: TypeAlias = str
-ConnectionName: TypeAlias = str
-DatasetIndex: TypeAlias = int
-QuantumIndex: TypeAlias = int
-DatastoreName: TypeAlias = str
-DimensionElementName: TypeAlias = str
-DataCoordinateValues: TypeAlias = list[DataIdValue]
+type TaskLabel = str
+type DatasetTypeName = str
+type ConnectionName = str
+type DatasetIndex = int
+type QuantumIndex = int
+type DatastoreName = str
+type DimensionElementName = str
+type DataCoordinateValues = list[DataIdValue]
 
 
-_T = TypeVar("_T", bound=pydantic.BaseModel)
-
 FORMAT_VERSION: int = 1
 """
 File format version number for new files.
@@ -453,8 +449,10 @@ class BaseQuantumGraphWriter:
         cdict_data: bytes | None = None,
         zstd_level: int = 10,
     ) -> Iterator[Self]:
-        uri = ResourcePath(uri)
+        uri = ResourcePath(uri, forceDirectory=False)
         address_writer = AddressWriter()
+        if uri.isLocal:
+            os.makedirs(uri.dirname().ospath, exist_ok=True)
         cdict = zstandard.ZstdCompressionDict(cdict_data) if cdict_data is not None else None
         compressor = zstandard.ZstdCompressor(level=zstd_level, dict_data=cdict)
         with uri.open(mode="wb") as stream:
@@ -595,9 +593,9 @@ class BaseQuantumGraphReader:
                     )
 
     @staticmethod
-    def _read_single_block_static(
-        name: str, model_type: type[_T], zf: zipfile.ZipFile, decompressor: Decompressor
-    ) -> _T:
+    def _read_single_block_static[T: pydantic.BaseModel](
+        name: str, model_type: type[T], zf: zipfile.ZipFile, decompressor: Decompressor
+    ) -> T:
         """Read a single compressed JSON block from a 'file' in a zip archive.
 
         Parameters
@@ -620,7 +618,7 @@ class BaseQuantumGraphReader:
         json_data = decompressor.decompress(compressed_data)
         return model_type.model_validate_json(json_data)
 
-    def _read_single_block(self, name: str, model_type: type[_T]) -> _T:
+    def _read_single_block[T: pydantic.BaseModel](self, name: str, model_type: type[T]) -> T:
         """Read a single compressed JSON block from a 'file' in a zip archive.
 
         Parameters

lsst/pipe/base/quantum_graph/_multiblock.py

@@ -43,25 +43,22 @@ import dataclasses
 import logging
 import tempfile
 import uuid
-from collections.abc import Iterator
+import zipfile
+from collections.abc import Iterator, Set
 from contextlib import contextmanager
 from io import BufferedReader, BytesIO
 from operator import attrgetter
-from typing import IO, TYPE_CHECKING, Protocol, TypeAlias, TypeVar
+from typing import IO, Protocol, TypeVar
 
 import pydantic
 
-if TYPE_CHECKING:
-    import zipfile
-
-
 _LOG = logging.getLogger(__name__)
 
 
 _T = TypeVar("_T", bound=pydantic.BaseModel)
 
 
-UUID_int: TypeAlias = int
+type UUID_int = int
 
 MAX_UUID_INT: UUID_int = 2**128
 
@@ -77,7 +74,7 @@ individual quanta (especially for execution).
 
 
 class Compressor(Protocol):
-    """A protocol for objects with a `compress` method that takes and returns
+    """A protocol for objects with a ``compress`` method that takes and returns
     `bytes`.
     """
 
@@ -208,11 +205,11 @@ class AddressWriter:
     addresses: list[dict[uuid.UUID, Address]] = dataclasses.field(default_factory=list)
     """Addresses to store with each UUID.
 
-    Every key in one of these dictionaries must have an entry in `indices`.
+    Every key in one of these dictionaries must have an entry in ``indices``.
     The converse is not true.
     """
 
-    def write(self, stream: IO[bytes], int_size: int) -> None:
+    def write(self, stream: IO[bytes], int_size: int, all_ids: Set[uuid.UUID] | None = None) -> None:
         """Write all addresses to a file-like object.
 
         Parameters
@@ -221,15 +218,17 @@ class AddressWriter:
             Binary file-like object.
         int_size : `int`
             Number of bytes to use for all integers.
+        all_ids : `~collections.abc.Set` [`uuid.UUID`], optional
+            Set of the union of all UUIDs in any dictionary from a call to
+            `get_all_ids`.
         """
-        indices: set[uuid.UUID] = set()
-        for address_map in self.addresses:
-            indices.update(address_map.keys())
+        if all_ids is None:
+            all_ids = self.get_all_ids()
         stream.write(int_size.to_bytes(1))
-        stream.write(len(indices).to_bytes(int_size))
+        stream.write(len(all_ids).to_bytes(int_size))
         stream.write(len(self.addresses).to_bytes(int_size))
         empty_address = Address()
-        for n, key in enumerate(sorted(indices, key=attrgetter("int"))):
+        for n, key in enumerate(sorted(all_ids, key=attrgetter("int"))):
             row = AddressRow(key, n, [m.get(key, empty_address) for m in self.addresses])
             _LOG.debug("Wrote address %s.", row)
             row.write(stream, int_size)
@@ -246,8 +245,25 @@ class AddressWriter:
         int_size : `int`
             Number of bytes to use for all integers.
         """
-        with zf.open(f"{name}.addr", mode="w") as stream:
-            self.write(stream, int_size=int_size)
+        all_ids = self.get_all_ids()
+        zip_info = zipfile.ZipInfo(f"{name}.addr")
+        row_size = AddressReader.compute_row_size(int_size, len(self.addresses))
+        zip_info.file_size = AddressReader.compute_header_size(int_size) + len(all_ids) * row_size
+        with zf.open(zip_info, mode="w") as stream:
+            self.write(stream, int_size=int_size, all_ids=all_ids)
+
+    def get_all_ids(self) -> Set[uuid.UUID]:
+        """Return all IDs used by any address dictionary.
+
+        Returns
+        -------
+        all_ids : `~collections.abc.Set` [`uuid.UUID`]
+            Set of all IDs.
+        """
+        all_ids: set[uuid.UUID] = set()
+        for address_map in self.addresses:
+            all_ids.update(address_map.keys())
+        return all_ids
 
 
 @dataclasses.dataclass
@@ -646,7 +662,7 @@ class MultiblockWriter:
         model : `pydantic.BaseModel`
             Model to convert to JSON and compress.
         compressor : `Compressor`
-            Object with a `compress` method that takes and returns `bytes`.
+            Object with a ``compress`` method that takes and returns `bytes`.
 
         Returns
         -------
@@ -743,7 +759,7 @@ class MultiblockReader:
         model_type : `type` [ `pydantic.BaseModel` ]
             Pydantic model to validate JSON with.
         decompressor : `Decompressor`
-            Object with a `decompress` method that takes and returns `bytes`.
+            Object with a ``decompress`` method that takes and returns `bytes`.
         int_size : `int`
             Number of bytes to use for all integers.
         page_size : `int`
@@ -793,7 +809,7 @@ class MultiblockReader:
         model_type : `type` [ `pydantic.BaseModel` ]
             Pydantic model to validate JSON with.
         decompressor : `Decompressor`
-            Object with a `decompress` method that takes and returns `bytes`.
+            Object with a ``decompress`` method that takes and returns `bytes`.
 
         Returns
         -------

lsst/pipe/base/quantum_graph/_predicted.py

@@ -49,7 +49,7 @@ import warnings
 from collections import defaultdict
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from contextlib import AbstractContextManager, contextmanager
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import networkx
 import networkx.algorithms.bipartite
@@ -110,10 +110,14 @@ if TYPE_CHECKING:
     from ..config import PipelineTaskConfig
     from ..graph import QgraphSummary, QuantumGraph
 
-_LOG = logging.getLogger(__name__)
+# Sphinx needs imports for type annotations of base class members.
+if "sphinx" in sys.modules:
+    import zipfile  # noqa: F401
+
+    from ._multiblock import AddressReader, Decompressor  # noqa: F401
 
 
-_T = TypeVar("_T", bound=pydantic.BaseModel)
+_LOG = logging.getLogger(__name__)
 
 
 class _PredictedThinQuantumModelV0(pydantic.BaseModel):