Coverage for python/lsst/pipe/base/quantum_graph/

1# This file is part of pipe_base.

3# Developed for the LSST Data Management System.

4# This product includes software developed by the LSST Project

5# (http://www.lsst.org).

6# See the COPYRIGHT file at the top-level directory of this distribution

7# for details of code ownership.

9# This software is dual licensed under the GNU General Public License and also

10# under a 3-clause BSD license. Recipients may choose which of these licenses

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,

12# respectively. If you choose the GPL option then the following text applies

13# (but note that there is still no warranty even if you opt for BSD instead):

14#

15# This program is free software: you can redistribute it and/or modify

16# it under the terms of the GNU General Public License as published by

17# the Free Software Foundation, either version 3 of the License, or

18# (at your option) any later version.

19#

20# This program is distributed in the hope that it will be useful,

21# but WITHOUT ANY WARRANTY; without even the implied warranty of

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

23# GNU General Public License for more details.

24#

25# You should have received a copy of the GNU General Public License

26# along with this program. If not, see <http://www.gnu.org/licenses/>.

28from __future__ import annotations

30__all__ = (

31 "Address",

32 "AddressReader",

33 "AddressRow",

34 "AddressWriter",

35 "Compressor",

36 "Decompressor",

37 "InvalidQuantumGraphFileError",

38 "MultiblockReader",

39 "MultiblockWriter",

40)

42import dataclasses

43import logging

44import tempfile

45import uuid

46from collections.abc import Iterator

47from contextlib import contextmanager

48from io import BufferedReader, BytesIO

49from operator import attrgetter

50from typing import IO, TYPE_CHECKING, Protocol, TypeAlias, TypeVar

52import pydantic

54if TYPE_CHECKING:

55 import zipfile

58_LOG = logging.getLogger(__name__)

61_T = TypeVar("_T", bound=pydantic.BaseModel)

64UUID_int: TypeAlias = int

66MAX_UUID_INT: UUID_int = 2**128

69DEFAULT_PAGE_SIZE: int = 5_000_000

70"""Default page size for reading chunks of quantum graph files.

72This is intended to be large enough to avoid any possibility of individual

73reads suffering from per-seek overheads, especially in network file access,

74while still being small enough to only minimally slow down tiny reads of

75individual quanta (especially for execution).

76"""

79class Compressor(Protocol):

80 """A protocol for objects with a `compress` method that takes and returns

81 `bytes`.

82 """

84 def compress(self, data: bytes) -> bytes:

85 """Compress the given data.

87 Parameters

88 ----------

89 data : `bytes`

90 Uncompressed data.

92 Returns

93 -------

94 compressed : `bytes`

95 Compressed data.

96 """

97 ...

100class Decompressor(Protocol):

101 """A protocol for objects with a `decompress` method that takes and returns

102 `bytes`.

103 """

104

105 def decompress(self, data: bytes) -> bytes:

106 """Decompress the given data.

107

108 Parameters

109 ----------

110 data : `bytes`

111 Compressed data.

112

113 Returns

114 -------

115 decompressed : `bytes`

116 Uncompressed data.

117 """

118 ...

119

120

121class InvalidQuantumGraphFileError(RuntimeError):

122 """An exception raised when a quantum graph file has internal

123 inconsistencies or does not actually appear to be a quantum graph file.

124 """

125

126

127@dataclasses.dataclass(slots=True)

128class Address:

129 """Struct that holds an address into a multi-block file."""

130

131 offset: int = 0

132 """Byte offset for the block."""

133

134 size: int = 0

135 """Size of the block.

136

137 This always includes the size of the tiny header that records the block

138 size. That header does not include the size of the header, so these sizes

139 differ by the ``int_size`` used to write the multi-block file.

140

141 A size of zero is used (with, by convention, an offset of zero) to indicate

142 an absent block.

143 """

144

145 def __str__(self) -> str:

146 return f"{self.offset:06}[{self.size:06}]"

147

148

149@dataclasses.dataclass(slots=True)

150class AddressRow:

151 """The in-memory representation of a single row in an address file."""

152

153 key: uuid.UUID

154 """Universally unique identifier for this row."""

155

156 index: int

157 """Monotonically increasing integer ID; unique within this file only."""

158

159 addresses: list[Address] = dataclasses.field(default_factory=list)

160 """Offsets and sizes into multi-block files."""

161

162 def write(self, stream: IO[bytes], int_size: int) -> None:

163 """Write this address row to a file-like object.

164

165 Parameters

166 ----------

167 stream : `typing.IO` [ `bytes` ]

168 Binary file-like object.

169 int_size : `int`

170 Number of bytes to use for all integers.

171 """

172 stream.write(self.key.bytes)

173 stream.write(self.index.to_bytes(int_size))

174 for address in self.addresses:

175 stream.write(address.offset.to_bytes(int_size))

176 stream.write(address.size.to_bytes(int_size))

177

178 @classmethod

179 def read(cls, stream: IO[bytes], n_addresses: int, int_size: int) -> AddressRow:

180 """Read this address row from a file-like object.

181

182 Parameters

183 ----------

184 stream : `typing.IO` [ `bytes` ]

185 Binary file-like object.

186 n_addresses : `int`

187 Number of addresses included in each row.

188 int_size : `int`

189 Number of bytes to use for all integers.

190 """

191 key = uuid.UUID(int=int.from_bytes(stream.read(16)))

192 index = int.from_bytes(stream.read(int_size))

193 row = AddressRow(key, index)

194 for _ in range(n_addresses):

195 offset = int.from_bytes(stream.read(int_size))

196 size = int.from_bytes(stream.read(int_size))

197 row.addresses.append(Address(offset, size))

198 return row

199

200 def __str__(self) -> str:

201 return f"{self.key} {self.index:06} {' '.join(str(a) for a in self.addresses)}"

202

203

204@dataclasses.dataclass

205class AddressWriter:

206 """A helper object for writing address files for multi-block files."""

207

208 indices: dict[uuid.UUID, int] = dataclasses.field(default_factory=dict)

209 """Mapping from UUID to internal integer ID.

210

211 The internal integer ID must always correspond to the index into the

212 sorted list of all UUIDs, but this `dict` need not be sorted itself.

213 """

214

215 addresses: list[dict[uuid.UUID, Address]] = dataclasses.field(default_factory=list)

216 """Addresses to store with each UUID.

217

218 Every key in one of these dictionaries must have an entry in `indices`.

219 The converse is not true.

220 """

221

222 def write(self, stream: IO[bytes], int_size: int) -> None:

223 """Write all addresses to a file-like object.

224

225 Parameters

226 ----------

227 stream : `typing.IO` [ `bytes` ]

228 Binary file-like object.

229 int_size : `int`

230 Number of bytes to use for all integers.

231 """

232 for n, address_map in enumerate(self.addresses):

233 if not self.indices.keys() >= address_map.keys():

234 raise AssertionError(

235 f"Logic bug in quantum graph I/O: address map {n} of {len(self.addresses)} has IDs "

236 f"{address_map.keys() - self.indices.keys()} not in the index map."

237 )

238 stream.write(int_size.to_bytes(1))

239 stream.write(len(self.indices).to_bytes(int_size))

240 stream.write(len(self.addresses).to_bytes(int_size))

241 empty_address = Address()

242 for key in sorted(self.indices.keys(), key=attrgetter("int")):

243 row = AddressRow(key, self.indices[key], [m.get(key, empty_address) for m in self.addresses])

244 _LOG.debug("Wrote address %s.", row)

245 row.write(stream, int_size)

246

247 def write_to_zip(self, zf: zipfile.ZipFile, name: str, int_size: int) -> None:

248 """Write all addresses to a file in a zip archive.

249

250 Parameters

251 ----------

252 zf : `zipfile.ZipFile`

253 Zip archive to add the file to.

254 name : `str`

255 Base name for the address file; an extension will be added.

256 int_size : `int`

257 Number of bytes to use for all integers.

258 """

259 with zf.open(f"{name}.addr", mode="w") as stream:

260 self.write(stream, int_size=int_size)

261

262

263@dataclasses.dataclass

264class AddressPage:

265 """A page of addresses in the `AddressReader`."""

266

267 file_offset: int

268 """Offset in bytes to this page from the beginning of the file."""

269

270 begin: int

271 """Index of the first row in this page."""

272

273 n_rows: int

274 """Number of rows in this page."""

275

276 read: bool = False

277 """Whether this page has already been read."""

278

279 @property

280 def end(self) -> int:

281 """One past the last row index in this page."""

282 return self.begin + self.n_rows

283

284

285@dataclasses.dataclass

286class PageBounds:

287 """A page index and the UUID interval that page covers."""

288

289 page_index: int

290 """Index into the page array."""

291

292 uuid_int_begin: UUID_int

293 """Integer representation of the smallest UUID in this page."""

294

295 uuid_int_end: UUID_int

296 """One larger than the integer representation of the largest UUID in this

297 page.

298 """

299

300 def __str__(self) -> str:

301 return f"{self.page_index} [{self.uuid_int_begin:x}:{self.uuid_int_end:x}]"

302

303

304@dataclasses.dataclass

305class AddressReader:

306 """A helper object for reading address files for multi-block files."""

307

308 stream: IO[bytes]

309 """Stream to read from."""

310

311 int_size: int

312 """Size of each integer in bytes."""

313

314 n_rows: int

315 """Number of rows in the file."""

316

317 n_addresses: int

318 """Number of addresses in each row."""

319

320 rows_per_page: int

321 """Number of addresses in each page."""

322

323 rows: dict[uuid.UUID, AddressRow] = dataclasses.field(default_factory=dict)

324 """Rows that have already been read."""

325

326 pages: list[AddressPage] = dataclasses.field(default_factory=list)

327 """Descriptions of the file offsets and integer row indexes of pages and

328 flags for whether they have been read already.

329 """

330

331 page_bounds: dict[int, PageBounds] = dataclasses.field(default_factory=dict)

332 """Mapping from page index to page boundary information."""

333

334 @classmethod

335 def from_stream(

336 cls, stream: IO[bytes], *, page_size: int, n_addresses: int, int_size: int

337 ) -> AddressReader:

338 """Construct from a stream by reading the header.

339

340 Parameters

341 ----------

342 stream : `typing.IO` [ `bytes` ]

343 File-like object to read from.

344 page_size : `int`

345 Approximate number of bytes to read at a time when searching for an

346 address.

347 n_addresses : `int`

348 Number of addresses to expect per row. This is checked against

349 the size embedded in the file.

350 int_size : `int`

351 Number of bytes to use for all integers. This is checked against

352 the size embedded in the file.

353 """

354 header_size = cls.compute_header_size(int_size)

355 row_size = cls.compute_row_size(int_size, n_addresses)

356 # Read the raw header page.

357 header_page_data = stream.read(header_size)

358 if len(header_page_data) < header_size:

359 raise InvalidQuantumGraphFileError("Address file unexpectedly truncated.")

360 # Interpret the raw header data and initialize the reader instance.

361 header_page_stream = BytesIO(header_page_data)

362 file_int_size = int.from_bytes(header_page_stream.read(1))

363 if file_int_size != int_size:

364 raise InvalidQuantumGraphFileError(

365 f"int size in address file ({file_int_size}) does not match int size in header ({int_size})."

366 )

367 n_rows = int.from_bytes(header_page_stream.read(int_size))

368 file_n_addresses = int.from_bytes(header_page_stream.read(int_size))

369 if file_n_addresses != n_addresses:

370 raise InvalidQuantumGraphFileError(

371 f"Incorrect number of addresses per row: expected {n_addresses}, got {file_n_addresses}."

372 )

373 rows_per_page = max(page_size // row_size, 1)

374 # Construct an instance.

375 self = cls(stream, int_size, n_rows, n_addresses, rows_per_page=rows_per_page)

376 # Calculate positions of each page of rows.

377 row_index = 0

378 file_offset = header_size

379 while row_index < n_rows:

380 self.pages.append(AddressPage(file_offset=file_offset, begin=row_index, n_rows=rows_per_page))

381 row_index += rows_per_page

382 file_offset += rows_per_page * row_size

383 if row_index != n_rows:

384 # Last page was too big.

385 self.pages[-1].n_rows -= row_index - n_rows

386 assert sum(p.n_rows for p in self.pages) == n_rows, "Bad logic setting page row counts."

387 return self

388

389 @classmethod

390 @contextmanager

391 def open_in_zip(

392 cls,

393 zf: zipfile.ZipFile,

394 name: str,

395 *,

396 page_size: int,

397 n_addresses: int,

398 int_size: int,

399 ) -> Iterator[AddressReader]:

400 """Make a reader for an address file in a zip archive.

401

402 Parameters

403 ----------

404 zf : `zipfile.ZipFile`

405 Zip archive to read the file from.

406 name : `str`

407 Base name for the address file; an extension will be added.

408 page_size : `int`

409 Approximate number of bytes to read at a time when searching for an

410 address.

411 n_addresses : `int`

412 Number of addresses to expect per row. This is checked against

413 the size embedded in the file.

414 int_size : `int`

415 Number of bytes to use for all integers. This is checked against

416 the size embedded in the file.

417

418 Returns

419 -------

420 reader : `contextlib.AbstractContextManager` [ `AddressReader` ]

421 Context manager that returns a reader when entered.

422 """

423 with zf.open(f"{name}.addr", mode="r") as stream:

424 yield cls.from_stream(stream, page_size=page_size, n_addresses=n_addresses, int_size=int_size)

425

426 @staticmethod

427 def compute_header_size(int_size: int) -> int:

428 """Return the size (in bytes) of the header of an address file.

429

430 Parameters

431 ----------

432 int_size : `int`

433 Size of each integer in bytes.

434

435 Returns

436 -------

437 size : `int`

438 Size of the header in bytes.

439 """

440 return (

441 1 # int_size

442 + int_size # number of rows

443 + int_size # number of addresses in each row

444 )

445

446 @staticmethod

447 def compute_row_size(int_size: int, n_addresses: int) -> int:

448 """Return the size (in bytes) of each row of an address file.

449

450 Parameters

451 ----------

452 int_size : `int`

453 Size of each integer in bytes.

454 n_addresses : `int`

455 Number of addresses in each row.

456

457 Returns

458 -------

459 size : `int`

460 Size of each row in bytes.

461 """

462 return (

463 16 # uuid

464 + int_size

465 * (

466 1 # index

467 + 2 * n_addresses

468 )

469 )

470

471 @property

472 def row_size(self) -> int:

473 """The size (in bytes) of each row of this address file."""

474 return self.compute_row_size(self.int_size, self.n_addresses)

475

476 def read_all(self) -> dict[uuid.UUID, AddressRow]:

477 """Read all addresses in the file.

478

479 Returns

480 -------

481 rows : `dict` [ `uuid.UUID`, `AddressRow` ]

482 Mapping of loaded address rows, keyed by UUID.

483 """

484 # Skip any pages from the beginning that have already been read; this

485 # nicely handles both the case where we already read everything (or

486 # there was nothing to read) while giving us a page with a file offset

487 # to start from.

488 for page in self.pages:

489 if not page.read:

490 break

491 else:

492 return self.rows

493 # Read the entire rest of the file into memory.

494 self.stream.seek(page.file_offset)

495 data = self.stream.read()

496 buffer = BytesIO(data)

497 # Shortcut out if we've already read everything, but don't bother

498 # optimizing previous partial reads.

499 while len(self.rows) < self.n_rows:

500 self._read_row(buffer)

501 # Delete all pages; they don't matter anymore, and that's easier than

502 # updating them to reflect the reads we've done.

503 self.pages.clear()

504 return self.rows

505

506 def find(self, key: uuid.UUID) -> AddressRow:

507 """Read the row for the given UUID or integer index.

508

509 Parameters

510 ----------

511 key : `uuid.UUID`

512 UUID to find.

513

514 Returns

515 -------

516 row : `AddressRow`

517 Addresses for the given UUID.

518 """

519 if (row := self.rows.get(key)) is not None:

520 return row

521 if self.n_rows == 0 or not self.pages:

522 raise LookupError(f"Address for {key} not found.")

523

524 # Use a binary search to find the page containing the target UUID.

525 left = 0

526 right = len(self.pages) - 1

527 while left <= right:

528 mid = left + ((right - left) // 2)

529 self._read_page(mid)

530 if (row := self.rows.get(key)) is not None:

531 return row

532 bounds = self.page_bounds[mid]

533 if key.int < bounds.uuid_int_begin:

534 right = mid - 1

535 elif key.int > bounds.uuid_int_end:

536 left = mid + 1

537 else:

538 # Should have been on this page, but it wasn't.

539 raise LookupError(f"Address for {key} not found.")

540

541 # Ran out of pages to search.

542 raise LookupError(f"Address for {key} not found.")

543

544 def _read_page(self, page_index: int, page_stream: BytesIO | None = None) -> bool:

545 page = self.pages[page_index]

546 if page.read:

547 return False

548 if page_stream is None:

549 self.stream.seek(page.file_offset)

550 page_stream = BytesIO(self.stream.read(page.n_rows * self.row_size))

551 row = self._read_row(page_stream)

552 uuid_int_begin = row.key.int

553 for _ in range(1, page.n_rows):

554 row = self._read_row(page_stream)

555 uuid_int_end = row.key.int + 1 # Python's loop scoping rules are actually useful here!

556 page.read = True

557 bounds = PageBounds(page_index=page_index, uuid_int_begin=uuid_int_begin, uuid_int_end=uuid_int_end)

558 self.page_bounds[page_index] = bounds

559 _LOG.debug("Read page %s with rows [%s:%s].", bounds, page.begin, page.end)

560 return True

561

562 def _read_row(self, page_stream: BytesIO) -> AddressRow:

563 row = AddressRow.read(page_stream, self.n_addresses, self.int_size)

564 self.rows[row.key] = row

565 _LOG.debug("Read address row %s.", row)

566 return row

567

568

569@dataclasses.dataclass

570class MultiblockWriter:

571 """A helper object for writing multi-block files."""

572

573 stream: IO[bytes]

574 """A binary file-like object to write to."""

575

576 int_size: int

577 """Number of bytes to use for all integers."""

578

579 file_size: int = 0

580 """Running size of the full file."""

581

582 addresses: dict[uuid.UUID, Address] = dataclasses.field(default_factory=dict)

583 """Running map of all addresses added to the file so far.

584

585 When the multi-block file is fully written, this is appended to the

586 `AddressWriter.addresses` to write the corresponding address file.

587 """

588

589 @classmethod

590 @contextmanager

591 def open_in_zip(

592 cls, zf: zipfile.ZipFile, name: str, int_size: int, use_tempfile: bool = False

593 ) -> Iterator[MultiblockWriter]:

594 """Open a writer for a file in a zip archive.

595

596 Parameters

597 ----------

598 zf : `zipfile.ZipFile`

599 Zip archive to add the file to.

600 name : `str`

601 Base name for the multi-block file; an extension will be added.

602 int_size : `int`

603 Number of bytes to use for all integers.

604 use_tempfile : `bool`, optional

605 If `True`, send writes to a temporary file and only add the file to

606 the zip archive when the context manager closes. This involves

607 more overall I/O, but it permits multiple multi-block files to be

608 open for writing in the same zip archive at once.

609

610 Returns

611 -------

612 writer : `contextlib.AbstractContextManager` [ `MultiblockWriter` ]

613 Context manager that returns a writer when entered.

614 """

615 filename = f"{name}.mb"

616 if use_tempfile:

617 with tempfile.NamedTemporaryFile(suffix=filename) as tmp:

618 yield MultiblockWriter(tmp, int_size)

619 tmp.flush()

620 zf.write(tmp.name, filename)

621 else:

622 with zf.open(f"{name}.mb", mode="w", force_zip64=True) as stream:

623 yield MultiblockWriter(stream, int_size)

624

625 def write_bytes(self, id: uuid.UUID, data: bytes) -> Address:

626 """Write raw bytes to the multi-block file.

627

628 Parameters

629 ----------

630 id : `uuid.UUID`

631 Unique ID of the object described by this block.

632 data : `bytes`

633 Data to store directly.

634

635 Returns

636 -------

637 address : `Address`

638 Address of the bytes just written.

639 """

640 assert id not in self.addresses, "Duplicate write to multi-block file detected."

641 self.stream.write(len(data).to_bytes(self.int_size))

642 self.stream.write(data)

643 block_size = len(data) + self.int_size

644 address = Address(offset=self.file_size, size=block_size)

645 self.file_size += block_size

646 self.addresses[id] = address

647 return address

648

649 def write_model(self, id: uuid.UUID, model: pydantic.BaseModel, compressor: Compressor) -> Address:

650 """Write raw bytes to the multi-block file.

651

652 Parameters

653 ----------

654 id : `uuid.UUID`

655 Unique ID of the object described by this block.

656 model : `pydantic.BaseModel`

657 Model to convert to JSON and compress.

658 compressor : `Compressor`

659 Object with a `compress` method that takes and returns `bytes`.

660

661 Returns

662 -------

663 address : `Address`

664 Address of the bytes just written.

665 """

666 json_data = model.model_dump_json().encode()

667 compressed_data = compressor.compress(json_data)

668 return self.write_bytes(id, compressed_data)

669

670

671@dataclasses.dataclass

672class MultiblockReader:

673 """A helper object for reader multi-block files."""

674

675 stream: IO[bytes]

676 """A binary file-like object to read from."""

677

678 int_size: int

679 """Number of bytes to use for all integers."""

680

681 @classmethod

682 @contextmanager

683 def open_in_zip(cls, zf: zipfile.ZipFile, name: str, *, int_size: int) -> Iterator[MultiblockReader]:

684 """Open a reader for a file in a zip archive.

685

686 Parameters

687 ----------

688 zf : `zipfile.ZipFile`

689 Zip archive to read the file from.

690 name : `str`

691 Base name for the multi-block file; an extension will be added.

692 int_size : `int`

693 Number of bytes to use for all integers.

694

695 Returns

696 -------

697 reader : `contextlib.AbstractContextManager` [ `MultiblockReader` ]

698 Context manager that returns a reader when entered.

699 """

700 with zf.open(f"{name}.mb", mode="r") as stream:

701 yield MultiblockReader(stream, int_size)

702

703 @classmethod

704 def read_all_bytes_in_zip(

705 cls, zf: zipfile.ZipFile, name: str, *, int_size: int, page_size: int

706 ) -> Iterator[bytes]:

707 """Iterate over all of the byte blocks in a file in a zip archive.

708

709 Parameters

710 ----------

711 zf : `zipfile.ZipFile`

712 Zip archive to read the file from.

713 name : `str`

714 Base name for the multi-block file; an extension will be added.

715 int_size : `int`

716 Number of bytes to use for all integers.

717 page_size : `int`

718 Approximate number of bytes to read at a time.

719

720 Returns

721 -------

722 byte_iter : `~collections.abc.Iterator` [ `bytes` ]

723 Iterator over blocks.

724 """

725 with zf.open(f"{name}.mb", mode="r") as zf_stream:

726 # The standard library typing of IO[bytes] tiers isn't consistent.

727 buffered_stream = BufferedReader(zf_stream) # type: ignore[type-var]

728 size_data = buffered_stream.read(int_size)

729 while size_data:

730 internal_size = int.from_bytes(size_data)

731 yield buffered_stream.read(internal_size)

732 size_data = buffered_stream.read(int_size)

733

734 @classmethod

735 def read_all_models_in_zip(

736 cls,

737 zf: zipfile.ZipFile,

738 name: str,

739 model_type: type[_T],

740 decompressor: Decompressor,

741 *,

742 int_size: int,

743 page_size: int,

744 ) -> Iterator[_T]:

745 """Iterate over all of the models in a file in a zip archive.

746

747 Parameters

748 ----------

749 zf : `zipfile.ZipFile`

750 Zip archive to read the file from.

751 name : `str`

752 Base name for the multi-block file; an extension will be added.

753 model_type : `type` [ `pydantic.BaseModel` ]

754 Pydantic model to validate JSON with.

755 decompressor : `Decompressor`

756 Object with a `decompress` method that takes and returns `bytes`.

757 int_size : `int`

758 Number of bytes to use for all integers.

759 page_size : `int`

760 Approximate number of bytes to read at a time.

761

762 Returns

763 -------

764 model_iter : `~collections.abc.Iterator` [ `pydantic.BaseModel` ]

765 Iterator over model instances.

766 """

767 for compressed_data in cls.read_all_bytes_in_zip(zf, name, int_size=int_size, page_size=page_size):

768 json_data = decompressor.decompress(compressed_data)

769 yield model_type.model_validate_json(json_data)

770

771 def read_bytes(self, address: Address) -> bytes | None:

772 """Read raw bytes from the multi-block file.

773

774 Parameters

775 ----------

776 address : `Address`

777 Offset and size of the data to read.

778

779 Returns

780 -------

781 data : `bytes` or `None`

782 Data read directly, or `None` if the address has zero size.

783 """

784 if not address.size:

785 return None

786 self.stream.seek(address.offset)

787 data = self.stream.read(address.size)

788 internal_size = int.from_bytes(data[: self.int_size])

789 data = data[self.int_size :]

790 if len(data) != internal_size:

791 raise InvalidQuantumGraphFileError(

792 f"Internal size {internal_size} does not match loaded data size {len(data)}."

793 )

794 return data

795

796 def read_model(self, address: Address, model_type: type[_T], decompressor: Decompressor) -> _T | None:

797 """Read a single compressed JSON block.

798

799 Parameters

800 ----------

801 address : `Address`

802 Size and offset of the block.

803 model_type : `type` [ `pydantic.BaseModel` ]

804 Pydantic model to validate JSON with.

805 decompressor : `Decompressor`

806 Object with a `decompress` method that takes and returns `bytes`.

807

808 Returns

809 -------

810 model : `pydantic.BaseModel`

811 Validated model.

812 """

813 compressed_data = self.read_bytes(address)

814 if compressed_data is None:

815 return None

816 json_data = decompressor.decompress(compressed_data)

817 return model_type.model_validate_json(json_data)

Coverage for python / lsst / pipe / base / quantum_graph / _multiblock.py: 39%

289 statements