Coverage for python/lsst/pipe/base/quantum_graph/_multiblock.py: 37%
311 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-13 01:39 -0800
1# This file is part of pipe_base.
2#
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.
8#
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 """
105 def decompress(self, data: bytes) -> bytes:
106 """Decompress the given data.
108 Parameters
109 ----------
110 data : `bytes`
111 Compressed data.
113 Returns
114 -------
115 decompressed : `bytes`
116 Uncompressed data.
117 """
118 ...
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 """
127@dataclasses.dataclass(slots=True)
128class Address:
129 """Struct that holds an address into a multi-block file."""
131 offset: int = 0
132 """Byte offset for the block."""
134 size: int = 0
135 """Size of the block.
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.
141 A size of zero is used (with, by convention, an offset of zero) to indicate
142 an absent block.
143 """
145 def __str__(self) -> str:
146 return f"{self.offset:06}[{self.size:06}]"
149@dataclasses.dataclass(slots=True)
150class AddressRow:
151 """The in-memory representation of a single row in an address file."""
153 key: uuid.UUID
154 """Universally unique identifier for this row."""
156 index: int
157 """Monotonically increasing integer ID; unique within this file only."""
159 addresses: list[Address] = dataclasses.field(default_factory=list)
160 """Offsets and sizes into multi-block files."""
162 def write(self, stream: IO[bytes], int_size: int) -> None:
163 """Write this address row to a file-like object.
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))
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.
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
200 def __str__(self) -> str:
201 return f"{self.key} {self.index:06} {' '.join(str(a) for a in self.addresses)}"
204@dataclasses.dataclass
205class AddressWriter:
206 """A helper object for writing address files for multi-block files."""
208 indices: dict[uuid.UUID, int] = dataclasses.field(default_factory=dict)
209 """Mapping from UUID to internal integer ID.
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 """
215 addresses: list[dict[uuid.UUID, Address]] = dataclasses.field(default_factory=list)
216 """Addresses to store with each UUID.
218 Every key in one of these dictionaries must have an entry in `indices`.
219 The converse is not true.
220 """
222 def write(self, stream: IO[bytes], int_size: int) -> None:
223 """Write all addresses to a file-like object.
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)
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.
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)
263@dataclasses.dataclass
264class AddressPage:
265 """A page of addresses in the `AddressReader`."""
267 file_offset: int
268 """Offset in bytes to this page from the beginning of the file."""
270 begin: int
271 """Index of the first row in this page."""
273 n_rows: int
274 """Number of rows in this page."""
276 read: bool = False
277 """Whether this page has already been read."""
279 @property
280 def end(self) -> int:
281 """One past the last row index in this page."""
282 return self.begin + self.n_rows
285@dataclasses.dataclass
286class PageBounds:
287 """A page index and the UUID interval that page covers."""
289 page_index: int
290 """Index into the page array."""
292 uuid_int_begin: UUID_int
293 """Integer representation of the smallest UUID in this page."""
295 uuid_int_end: UUID_int
296 """One larger than the integer representation of the largest UUID in this
297 page.
298 """
300 def __str__(self) -> str:
301 return f"{self.page_index} [{self.uuid_int_begin:x}:{self.uuid_int_end:x}]"
304@dataclasses.dataclass
305class AddressReader:
306 """A helper object for reading address files for multi-block files."""
308 stream: IO[bytes]
309 """Stream to read from."""
311 int_size: int
312 """Size of each integer in bytes."""
314 n_rows: int
315 """Number of rows in the file."""
317 n_addresses: int
318 """Number of addresses in each row."""
320 rows_per_page: int
321 """Number of addresses in each page."""
323 rows: dict[uuid.UUID, AddressRow] = dataclasses.field(default_factory=dict)
324 """Rows that have already been read."""
326 rows_by_index: dict[int, AddressRow] = dataclasses.field(default_factory=dict)
327 """Rows that have already been read, keyed by integer index."""
329 pages: list[AddressPage] = dataclasses.field(default_factory=list)
330 page_bounds: dict[int, PageBounds] = dataclasses.field(default_factory=dict)
331 """Mapping from page index to page boundary information."""
333 @classmethod
334 def from_stream(
335 cls, stream: IO[bytes], *, page_size: int, n_addresses: int, int_size: int
336 ) -> AddressReader:
337 """Construct from a stream by reading the header.
339 Parameters
340 ----------
341 stream : `typing.IO` [ `bytes` ]
342 File-like object to read from.
343 page_size : `int`
344 Approximate number of bytes to read at a time when searching for an
345 address.
346 n_addresses : `int`
347 Number of addresses to expect per row. This is checked against
348 the size embedded in the file.
349 int_size : `int`
350 Number of bytes to use for all integers. This is checked against
351 the size embedded in the file.
352 """
353 header_size = cls.compute_header_size(int_size)
354 row_size = cls.compute_row_size(int_size, n_addresses)
355 # Read the raw header page.
356 header_page_data = stream.read(header_size)
357 if len(header_page_data) < header_size:
358 raise InvalidQuantumGraphFileError("Address file unexpectedly truncated.")
359 # Interpret the raw header data and initialize the reader instance.
360 header_page_stream = BytesIO(header_page_data)
361 file_int_size = int.from_bytes(header_page_stream.read(1))
362 if file_int_size != int_size:
363 raise InvalidQuantumGraphFileError(
364 f"int size in address file ({file_int_size}) does not match int size in header ({int_size})."
365 )
366 n_rows = int.from_bytes(header_page_stream.read(int_size))
367 file_n_addresses = int.from_bytes(header_page_stream.read(int_size))
368 if file_n_addresses != n_addresses:
369 raise InvalidQuantumGraphFileError(
370 f"Incorrect number of addresses per row: expected {n_addresses}, got {file_n_addresses}."
371 )
372 rows_per_page = max(page_size // row_size, 1)
373 # Construct an instance.
374 self = cls(stream, int_size, n_rows, n_addresses, rows_per_page=rows_per_page)
375 # Calculate positions of each page of rows.
376 row_index = 0
377 file_offset = header_size
378 while row_index < n_rows:
379 self.pages.append(AddressPage(file_offset=file_offset, begin=row_index, n_rows=rows_per_page))
380 row_index += rows_per_page
381 file_offset += rows_per_page * row_size
382 if row_index != n_rows:
383 # Last page was too big.
384 self.pages[-1].n_rows -= row_index - n_rows
385 assert sum(p.n_rows for p in self.pages) == n_rows, "Bad logic setting page row counts."
386 return self
388 @classmethod
389 @contextmanager
390 def open_in_zip(
391 cls,
392 zf: zipfile.ZipFile,
393 name: str,
394 *,
395 page_size: int,
396 n_addresses: int,
397 int_size: int,
398 ) -> Iterator[AddressReader]:
399 """Make a reader for an address file in a zip archive.
401 Parameters
402 ----------
403 zf : `zipfile.ZipFile`
404 Zip archive to read the file from.
405 name : `str`
406 Base name for the address file; an extension will be added.
407 page_size : `int`
408 Approximate number of bytes to read at a time when searching for an
409 address.
410 n_addresses : `int`
411 Number of addresses to expect per row. This is checked against
412 the size embedded in the file.
413 int_size : `int`
414 Number of bytes to use for all integers. This is checked against
415 the size embedded in the file.
417 Returns
418 -------
419 reader : `contextlib.AbstractContextManager` [ `AddressReader` ]
420 Context manager that returns a reader when entered.
421 """
422 with zf.open(f"{name}.addr", mode="r") as stream:
423 yield cls.from_stream(stream, page_size=page_size, n_addresses=n_addresses, int_size=int_size)
425 @staticmethod
426 def compute_header_size(int_size: int) -> int:
427 """Return the size (in bytes) of the header of an address file.
429 Parameters
430 ----------
431 int_size : `int`
432 Size of each integer in bytes.
434 Returns
435 -------
436 size : `int`
437 Size of the header in bytes.
438 """
439 return (
440 1 # int_size
441 + int_size # number of rows
442 + int_size # number of addresses in each row
443 )
445 @staticmethod
446 def compute_row_size(int_size: int, n_addresses: int) -> int:
447 """Return the size (in bytes) of each row of an address file.
449 Parameters
450 ----------
451 int_size : `int`
452 Size of each integer in bytes.
453 n_addresses : `int`
454 Number of addresses in each row.
456 Returns
457 -------
458 size : `int`
459 Size of each row in bytes.
460 """
461 return (
462 16 # uuid
463 + int_size
464 * (
465 1 # index
466 + 2 * n_addresses
467 )
468 )
470 @property
471 def row_size(self) -> int:
472 """The size (in bytes) of each row of this address file."""
473 return self.compute_row_size(self.int_size, self.n_addresses)
475 def read_all(self) -> dict[uuid.UUID, AddressRow]:
476 """Read all addresses in the file.
478 Returns
479 -------
480 rows : `dict` [ `uuid.UUID`, `AddressRow` ]
481 Mapping of loaded address rows, keyed by UUID.
482 """
483 # Skip any pages from the beginning that have already been read; this
484 # nicely handles both the case where we already read everything (or
485 # there was nothing to read) while giving us a page with a file offset
486 # to start from.
487 for page in self.pages:
488 if not page.read:
489 break
490 else:
491 return self.rows
492 # Read the entire rest of the file into memory.
493 self.stream.seek(page.file_offset)
494 data = self.stream.read()
495 buffer = BytesIO(data)
496 # Shortcut out if we've already read everything, but don't bother
497 # optimizing previous partial reads.
498 while len(self.rows) < self.n_rows:
499 self._read_row(buffer)
500 # Delete all pages; they don't matter anymore, and that's easier than
501 # updating them to reflect the reads we've done.
502 self.pages.clear()
503 return self.rows
505 def find(self, key: uuid.UUID | int) -> AddressRow:
506 """Read the row for the given UUID or integer index.
508 Parameters
509 ----------
510 key : `uuid.UUID` or `int`
511 UUID or integer index to find.
513 Returns
514 -------
515 row : `AddressRow`
516 Addresses for the given UUID.
517 """
518 match key:
519 case uuid.UUID():
520 return self._find_uuid(key)
521 case int():
522 return self._find_index(key)
523 case _:
524 raise TypeError(f"Invalid argument: {key}.")
526 def _find_uuid(self, target: uuid.UUID) -> AddressRow:
527 if (row := self.rows.get(target)) is not None:
528 return row
529 if self.n_rows == 0 or not self.pages:
530 raise LookupError(f"Address for {target} not found.")
532 # Use a binary search to find the page containing the target UUID.
533 left = 0
534 right = len(self.pages) - 1
535 while left <= right:
536 mid = left + ((right - left) // 2)
537 self._read_page(mid)
538 if (row := self.rows.get(target)) is not None:
539 return row
540 bounds = self.page_bounds[mid]
541 if target.int < bounds.uuid_int_begin:
542 right = mid - 1
543 elif target.int > bounds.uuid_int_end:
544 left = mid + 1
545 else:
546 # Should have been on this page, but it wasn't.
547 raise LookupError(f"Address for {target} not found.")
549 # Ran out of pages to search.
550 raise LookupError(f"Address for {target} not found.")
552 def _find_index(self, target: int) -> AddressRow:
553 # First shortcut if we've already loaded this row.
554 if (row := self.rows_by_index.get(target)) is not None:
555 return row
556 if target < 0 or target >= self.n_rows:
557 raise LookupError(f"Address for index {target} not found.")
558 # Since all indexes should be present, we can predict the right page
559 # exactly.
560 page_index = target // self.rows_per_page
561 self._read_page(page_index)
562 try:
563 return self.rows_by_index[target]
564 except KeyError:
565 _LOG.debug("Index find failed: %s should have been in page %s.", target, page_index)
566 raise LookupError(f"Address for {target} not found.") from None
568 def _read_page(self, page_index: int, page_stream: BytesIO | None = None) -> bool:
569 page = self.pages[page_index]
570 if page.read:
571 return False
572 if page_stream is None:
573 self.stream.seek(page.file_offset)
574 page_stream = BytesIO(self.stream.read(page.n_rows * self.row_size))
575 row = self._read_row(page_stream)
576 uuid_int_begin = row.key.int
577 for _ in range(1, page.n_rows):
578 row = self._read_row(page_stream)
579 uuid_int_end = row.key.int + 1 # Python's loop scoping rules are actually useful here!
580 page.read = True
581 bounds = PageBounds(page_index=page_index, uuid_int_begin=uuid_int_begin, uuid_int_end=uuid_int_end)
582 self.page_bounds[page_index] = bounds
583 _LOG.debug("Read page %s with rows [%s:%s].", bounds, page.begin, page.end)
584 return True
586 def _read_row(self, page_stream: BytesIO) -> AddressRow:
587 row = AddressRow.read(page_stream, self.n_addresses, self.int_size)
588 self.rows[row.key] = row
589 self.rows_by_index[row.index] = row
590 _LOG.debug("Read address row %s.", row)
591 return row
594@dataclasses.dataclass
595class MultiblockWriter:
596 """A helper object for writing multi-block files."""
598 stream: IO[bytes]
599 """A binary file-like object to write to."""
601 int_size: int
602 """Number of bytes to use for all integers."""
604 file_size: int = 0
605 """Running size of the full file."""
607 addresses: dict[uuid.UUID, Address] = dataclasses.field(default_factory=dict)
608 """Running map of all addresses added to the file so far.
610 When the multi-block file is fully written, this is appended to the
611 `AddressWriter.addresses` to write the corresponding address file.
612 """
614 @classmethod
615 @contextmanager
616 def open_in_zip(
617 cls, zf: zipfile.ZipFile, name: str, int_size: int, use_tempfile: bool = False
618 ) -> Iterator[MultiblockWriter]:
619 """Open a writer for a file in a zip archive.
621 Parameters
622 ----------
623 zf : `zipfile.ZipFile`
624 Zip archive to add the file to.
625 name : `str`
626 Base name for the multi-block file; an extension will be added.
627 int_size : `int`
628 Number of bytes to use for all integers.
629 use_tempfile : `bool`, optional
630 If `True`, send writes to a temporary file and only add the file to
631 the zip archive when the context manager closes. This involves
632 more overall I/O, but it permits multiple multi-block files to be
633 open for writing in the same zip archive at once.
635 Returns
636 -------
637 writer : `contextlib.AbstractContextManager` [ `MultiblockWriter` ]
638 Context manager that returns a writer when entered.
639 """
640 filename = f"{name}.mb"
641 if use_tempfile:
642 with tempfile.NamedTemporaryFile(suffix=filename) as tmp:
643 yield MultiblockWriter(tmp, int_size)
644 tmp.flush()
645 zf.write(tmp.name, filename)
646 else:
647 with zf.open(f"{name}.mb", mode="w", force_zip64=True) as stream:
648 yield MultiblockWriter(stream, int_size)
650 def write_bytes(self, id: uuid.UUID, data: bytes) -> Address:
651 """Write raw bytes to the multi-block file.
653 Parameters
654 ----------
655 id : `uuid.UUID`
656 Unique ID of the object described by this block.
657 data : `bytes`
658 Data to store directly.
660 Returns
661 -------
662 address : `Address`
663 Address of the bytes just written.
664 """
665 assert id not in self.addresses, "Duplicate write to multi-block file detected."
666 self.stream.write(len(data).to_bytes(self.int_size))
667 self.stream.write(data)
668 block_size = len(data) + self.int_size
669 address = Address(offset=self.file_size, size=block_size)
670 self.file_size += block_size
671 self.addresses[id] = address
672 return address
674 def write_model(self, id: uuid.UUID, model: pydantic.BaseModel, compressor: Compressor) -> Address:
675 """Write raw bytes to the multi-block file.
677 Parameters
678 ----------
679 id : `uuid.UUID`
680 Unique ID of the object described by this block.
681 model : `pydantic.BaseModel`
682 Model to convert to JSON and compress.
683 compressor : `Compressor`
684 Object with a `compress` method that takes and returns `bytes`.
686 Returns
687 -------
688 address : `Address`
689 Address of the bytes just written.
690 """
691 json_data = model.model_dump_json().encode()
692 compressed_data = compressor.compress(json_data)
693 return self.write_bytes(id, compressed_data)
696@dataclasses.dataclass
697class MultiblockReader:
698 """A helper object for reader multi-block files."""
700 stream: IO[bytes]
701 """A binary file-like object to read from."""
703 int_size: int
704 """Number of bytes to use for all integers."""
706 @classmethod
707 @contextmanager
708 def open_in_zip(cls, zf: zipfile.ZipFile, name: str, *, int_size: int) -> Iterator[MultiblockReader]:
709 """Open a reader for a file in a zip archive.
711 Parameters
712 ----------
713 zf : `zipfile.ZipFile`
714 Zip archive to read the file from.
715 name : `str`
716 Base name for the multi-block file; an extension will be added.
717 int_size : `int`
718 Number of bytes to use for all integers.
720 Returns
721 -------
722 reader : `contextlib.AbstractContextManager` [ `MultiblockReader` ]
723 Context manager that returns a reader when entered.
724 """
725 with zf.open(f"{name}.mb", mode="r") as stream:
726 yield MultiblockReader(stream, int_size)
728 @classmethod
729 def read_all_bytes_in_zip(
730 cls, zf: zipfile.ZipFile, name: str, *, int_size: int, page_size: int
731 ) -> Iterator[bytes]:
732 """Iterate over all of the byte blocks in a file in a zip archive.
734 Parameters
735 ----------
736 zf : `zipfile.ZipFile`
737 Zip archive to read the file from.
738 name : `str`
739 Base name for the multi-block file; an extension will be added.
740 int_size : `int`
741 Number of bytes to use for all integers.
742 page_size : `int`
743 Approximate number of bytes to read at a time.
745 Returns
746 -------
747 byte_iter : `~collections.abc.Iterator` [ `bytes` ]
748 Iterator over blocks.
749 """
750 with zf.open(f"{name}.mb", mode="r") as zf_stream:
751 # The standard library typing of IO[bytes] tiers isn't consistent.
752 buffered_stream = BufferedReader(zf_stream) # type: ignore[type-var]
753 size_data = buffered_stream.read(int_size)
754 while size_data:
755 internal_size = int.from_bytes(size_data)
756 yield buffered_stream.read(internal_size)
757 size_data = buffered_stream.read(int_size)
759 @classmethod
760 def read_all_models_in_zip(
761 cls,
762 zf: zipfile.ZipFile,
763 name: str,
764 model_type: type[_T],
765 decompressor: Decompressor,
766 *,
767 int_size: int,
768 page_size: int,
769 ) -> Iterator[_T]:
770 """Iterate over all of the models in a file in a zip archive.
772 Parameters
773 ----------
774 zf : `zipfile.ZipFile`
775 Zip archive to read the file from.
776 name : `str`
777 Base name for the multi-block file; an extension will be added.
778 model_type : `type` [ `pydantic.BaseModel` ]
779 Pydantic model to validate JSON with.
780 decompressor : `Decompressor`
781 Object with a `decompress` method that takes and returns `bytes`.
782 int_size : `int`
783 Number of bytes to use for all integers.
784 page_size : `int`
785 Approximate number of bytes to read at a time.
787 Returns
788 -------
789 model_iter : `~collections.abc.Iterator` [ `pydantic.BaseModel` ]
790 Iterator over model instances.
791 """
792 for compressed_data in cls.read_all_bytes_in_zip(zf, name, int_size=int_size, page_size=page_size):
793 json_data = decompressor.decompress(compressed_data)
794 yield model_type.model_validate_json(json_data)
796 def read_bytes(self, address: Address) -> bytes | None:
797 """Read raw bytes from the multi-block file.
799 Parameters
800 ----------
801 address : `Address`
802 Offset and size of the data to read.
804 Returns
805 -------
806 data : `bytes` or `None`
807 Data read directly, or `None` if the address has zero size.
808 """
809 if not address.size:
810 return None
811 self.stream.seek(address.offset)
812 data = self.stream.read(address.size)
813 internal_size = int.from_bytes(data[: self.int_size])
814 data = data[self.int_size :]
815 if len(data) != internal_size:
816 raise InvalidQuantumGraphFileError(
817 f"Internal size {internal_size} does not match loaded data size {len(data)}."
818 )
819 return data
821 def read_model(self, address: Address, model_type: type[_T], decompressor: Decompressor) -> _T | None:
822 """Read a single compressed JSON block.
824 Parameters
825 ----------
826 address : `Address`
827 Size and offset of the block.
828 model_type : `type` [ `pydantic.BaseModel` ]
829 Pydantic model to validate JSON with.
830 decompressor : `Decompressor`
831 Object with a `decompress` method that takes and returns `bytes`.
833 Returns
834 -------
835 model : `pydantic.BaseModel`
836 Validated model.
837 """
838 compressed_data = self.read_bytes(address)
839 if compressed_data is None:
840 return None
841 json_data = decompressor.decompress(compressed_data)
842 return model_type.model_validate_json(json_data)