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

288 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-20 01:37 -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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ( 

31 "Address", 

32 "AddressReader", 

33 "AddressRow", 

34 "AddressWriter", 

35 "Compressor", 

36 "Decompressor", 

37 "InvalidQuantumGraphFileError", 

38 "MultiblockReader", 

39 "MultiblockWriter", 

40) 

41 

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 

51 

52import pydantic 

53 

54if TYPE_CHECKING: 

55 import zipfile 

56 

57 

58_LOG = logging.getLogger(__name__) 

59 

60 

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

62 

63 

64UUID_int: TypeAlias = int 

65 

66MAX_UUID_INT: UUID_int = 2**128 

67 

68 

69DEFAULT_PAGE_SIZE: int = 5_000_000 

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

71 

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""" 

77 

78 

79class Compressor(Protocol): 

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

81 `bytes`. 

82 """ 

83 

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

85 """Compress the given data. 

86 

87 Parameters 

88 ---------- 

89 data : `bytes` 

90 Uncompressed data. 

91 

92 Returns 

93 ------- 

94 compressed : `bytes` 

95 Compressed data. 

96 """ 

97 ... 

98 

99 

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 addresses: list[dict[uuid.UUID, Address]] = dataclasses.field(default_factory=list) 

209 """Addresses to store with each UUID. 

210 

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

212 The converse is not true. 

213 """ 

214 

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

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

217 

218 Parameters 

219 ---------- 

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

221 Binary file-like object. 

222 int_size : `int` 

223 Number of bytes to use for all integers. 

224 """ 

225 indices: set[uuid.UUID] = set() 

226 for address_map in self.addresses: 

227 indices.update(address_map.keys()) 

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

229 stream.write(len(indices).to_bytes(int_size)) 

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

231 empty_address = Address() 

232 for n, key in enumerate(sorted(indices, key=attrgetter("int"))): 

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

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

235 row.write(stream, int_size) 

236 

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

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

239 

240 Parameters 

241 ---------- 

242 zf : `zipfile.ZipFile` 

243 Zip archive to add the file to. 

244 name : `str` 

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

246 int_size : `int` 

247 Number of bytes to use for all integers. 

248 """ 

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

250 self.write(stream, int_size=int_size) 

251 

252 

253@dataclasses.dataclass 

254class AddressPage: 

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

256 

257 file_offset: int 

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

259 

260 begin: int 

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

262 

263 n_rows: int 

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

265 

266 read: bool = False 

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

268 

269 @property 

270 def end(self) -> int: 

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

272 return self.begin + self.n_rows 

273 

274 

275@dataclasses.dataclass 

276class PageBounds: 

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

278 

279 page_index: int 

280 """Index into the page array.""" 

281 

282 uuid_int_begin: UUID_int 

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

284 

285 uuid_int_end: UUID_int 

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

287 page. 

288 """ 

289 

290 def __str__(self) -> str: 

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

292 

293 

294@dataclasses.dataclass 

295class AddressReader: 

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

297 

298 stream: IO[bytes] 

299 """Stream to read from.""" 

300 

301 int_size: int 

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

303 

304 n_rows: int 

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

306 

307 n_addresses: int 

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

309 

310 rows_per_page: int 

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

312 

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

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

315 

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

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

318 flags for whether they have been read already. 

319 """ 

320 

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

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

323 

324 @classmethod 

325 def from_stream( 

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

327 ) -> AddressReader: 

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

329 

330 Parameters 

331 ---------- 

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

333 File-like object to read from. 

334 page_size : `int` 

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

336 address. 

337 n_addresses : `int` 

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

339 the size embedded in the file. 

340 int_size : `int` 

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

342 the size embedded in the file. 

343 """ 

344 header_size = cls.compute_header_size(int_size) 

345 row_size = cls.compute_row_size(int_size, n_addresses) 

346 # Read the raw header page. 

347 header_page_data = stream.read(header_size) 

348 if len(header_page_data) < header_size: 

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

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

351 header_page_stream = BytesIO(header_page_data) 

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

353 if file_int_size != int_size: 

354 raise InvalidQuantumGraphFileError( 

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

356 ) 

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

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

359 if file_n_addresses != n_addresses: 

360 raise InvalidQuantumGraphFileError( 

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

362 ) 

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

364 # Construct an instance. 

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

366 # Calculate positions of each page of rows. 

367 row_index = 0 

368 file_offset = header_size 

369 while row_index < n_rows: 

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

371 row_index += rows_per_page 

372 file_offset += rows_per_page * row_size 

373 if row_index != n_rows: 

374 # Last page was too big. 

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

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

377 return self 

378 

379 @classmethod 

380 @contextmanager 

381 def open_in_zip( 

382 cls, 

383 zf: zipfile.ZipFile, 

384 name: str, 

385 *, 

386 page_size: int, 

387 n_addresses: int, 

388 int_size: int, 

389 ) -> Iterator[AddressReader]: 

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

391 

392 Parameters 

393 ---------- 

394 zf : `zipfile.ZipFile` 

395 Zip archive to read the file from. 

396 name : `str` 

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

398 page_size : `int` 

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

400 address. 

401 n_addresses : `int` 

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

403 the size embedded in the file. 

404 int_size : `int` 

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

406 the size embedded in the file. 

407 

408 Returns 

409 ------- 

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

411 Context manager that returns a reader when entered. 

412 """ 

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

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

415 

416 @staticmethod 

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

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

419 

420 Parameters 

421 ---------- 

422 int_size : `int` 

423 Size of each integer in bytes. 

424 

425 Returns 

426 ------- 

427 size : `int` 

428 Size of the header in bytes. 

429 """ 

430 return ( 

431 1 # int_size 

432 + int_size # number of rows 

433 + int_size # number of addresses in each row 

434 ) 

435 

436 @staticmethod 

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

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

439 

440 Parameters 

441 ---------- 

442 int_size : `int` 

443 Size of each integer in bytes. 

444 n_addresses : `int` 

445 Number of addresses in each row. 

446 

447 Returns 

448 ------- 

449 size : `int` 

450 Size of each row in bytes. 

451 """ 

452 return ( 

453 16 # uuid 

454 + int_size 

455 * ( 

456 1 # index 

457 + 2 * n_addresses 

458 ) 

459 ) 

460 

461 @property 

462 def row_size(self) -> int: 

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

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

465 

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

467 """Read all addresses in the file. 

468 

469 Returns 

470 ------- 

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

472 Mapping of loaded address rows, keyed by UUID. 

473 """ 

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

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

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

477 # to start from. 

478 for page in self.pages: 

479 if not page.read: 

480 break 

481 else: 

482 return self.rows 

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

484 self.stream.seek(page.file_offset) 

485 data = self.stream.read() 

486 buffer = BytesIO(data) 

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

488 # optimizing previous partial reads. 

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

490 self._read_row(buffer) 

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

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

493 self.pages.clear() 

494 return self.rows 

495 

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

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

498 

499 Parameters 

500 ---------- 

501 key : `uuid.UUID` 

502 UUID to find. 

503 

504 Returns 

505 ------- 

506 row : `AddressRow` 

507 Addresses for the given UUID. 

508 """ 

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

510 return row 

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

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

513 

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

515 left = 0 

516 right = len(self.pages) - 1 

517 while left <= right: 

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

519 self._read_page(mid) 

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

521 return row 

522 bounds = self.page_bounds[mid] 

523 if key.int < bounds.uuid_int_begin: 

524 right = mid - 1 

525 elif key.int > bounds.uuid_int_end: 

526 left = mid + 1 

527 else: 

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

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

530 

531 # Ran out of pages to search. 

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

533 

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

535 page = self.pages[page_index] 

536 if page.read: 

537 return False 

538 if page_stream is None: 

539 self.stream.seek(page.file_offset) 

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

541 row = self._read_row(page_stream) 

542 uuid_int_begin = row.key.int 

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

544 row = self._read_row(page_stream) 

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

546 page.read = True 

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

548 self.page_bounds[page_index] = bounds 

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

550 return True 

551 

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

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

554 self.rows[row.key] = row 

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

556 return row 

557 

558 

559@dataclasses.dataclass 

560class MultiblockWriter: 

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

562 

563 stream: IO[bytes] 

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

565 

566 int_size: int 

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

568 

569 file_size: int = 0 

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

571 

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

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

574 

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

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

577 """ 

578 

579 @classmethod 

580 @contextmanager 

581 def open_in_zip( 

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

583 ) -> Iterator[MultiblockWriter]: 

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

585 

586 Parameters 

587 ---------- 

588 zf : `zipfile.ZipFile` 

589 Zip archive to add the file to. 

590 name : `str` 

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

592 int_size : `int` 

593 Number of bytes to use for all integers. 

594 use_tempfile : `bool`, optional 

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

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

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

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

599 

600 Returns 

601 ------- 

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

603 Context manager that returns a writer when entered. 

604 """ 

605 filename = f"{name}.mb" 

606 if use_tempfile: 

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

608 yield MultiblockWriter(tmp, int_size) 

609 tmp.flush() 

610 zf.write(tmp.name, filename) 

611 else: 

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

613 yield MultiblockWriter(stream, int_size) 

614 

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

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

617 

618 Parameters 

619 ---------- 

620 id : `uuid.UUID` 

621 Unique ID of the object described by this block. 

622 data : `bytes` 

623 Data to store directly. 

624 

625 Returns 

626 ------- 

627 address : `Address` 

628 Address of the bytes just written. 

629 """ 

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

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

632 self.stream.write(data) 

633 block_size = len(data) + self.int_size 

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

635 self.file_size += block_size 

636 self.addresses[id] = address 

637 return address 

638 

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

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

641 

642 Parameters 

643 ---------- 

644 id : `uuid.UUID` 

645 Unique ID of the object described by this block. 

646 model : `pydantic.BaseModel` 

647 Model to convert to JSON and compress. 

648 compressor : `Compressor` 

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

650 

651 Returns 

652 ------- 

653 address : `Address` 

654 Address of the bytes just written. 

655 """ 

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

657 compressed_data = compressor.compress(json_data) 

658 return self.write_bytes(id, compressed_data) 

659 

660 

661@dataclasses.dataclass 

662class MultiblockReader: 

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

664 

665 stream: IO[bytes] 

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

667 

668 int_size: int 

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

670 

671 @classmethod 

672 @contextmanager 

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

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

675 

676 Parameters 

677 ---------- 

678 zf : `zipfile.ZipFile` 

679 Zip archive to read the file from. 

680 name : `str` 

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

682 int_size : `int` 

683 Number of bytes to use for all integers. 

684 

685 Returns 

686 ------- 

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

688 Context manager that returns a reader when entered. 

689 """ 

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

691 yield MultiblockReader(stream, int_size) 

692 

693 @classmethod 

694 def read_all_bytes_in_zip( 

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

696 ) -> Iterator[bytes]: 

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

698 

699 Parameters 

700 ---------- 

701 zf : `zipfile.ZipFile` 

702 Zip archive to read the file from. 

703 name : `str` 

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

705 int_size : `int` 

706 Number of bytes to use for all integers. 

707 page_size : `int` 

708 Approximate number of bytes to read at a time. 

709 

710 Returns 

711 ------- 

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

713 Iterator over blocks. 

714 """ 

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

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

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

718 size_data = buffered_stream.read(int_size) 

719 while size_data: 

720 internal_size = int.from_bytes(size_data) 

721 yield buffered_stream.read(internal_size) 

722 size_data = buffered_stream.read(int_size) 

723 

724 @classmethod 

725 def read_all_models_in_zip( 

726 cls, 

727 zf: zipfile.ZipFile, 

728 name: str, 

729 model_type: type[_T], 

730 decompressor: Decompressor, 

731 *, 

732 int_size: int, 

733 page_size: int, 

734 ) -> Iterator[_T]: 

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

736 

737 Parameters 

738 ---------- 

739 zf : `zipfile.ZipFile` 

740 Zip archive to read the file from. 

741 name : `str` 

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

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

744 Pydantic model to validate JSON with. 

745 decompressor : `Decompressor` 

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

747 int_size : `int` 

748 Number of bytes to use for all integers. 

749 page_size : `int` 

750 Approximate number of bytes to read at a time. 

751 

752 Returns 

753 ------- 

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

755 Iterator over model instances. 

756 """ 

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

758 json_data = decompressor.decompress(compressed_data) 

759 yield model_type.model_validate_json(json_data) 

760 

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

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

763 

764 Parameters 

765 ---------- 

766 address : `Address` 

767 Offset and size of the data to read. 

768 

769 Returns 

770 ------- 

771 data : `bytes` or `None` 

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

773 """ 

774 if not address.size: 

775 return None 

776 self.stream.seek(address.offset) 

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

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

779 data = data[self.int_size :] 

780 if len(data) != internal_size: 

781 raise InvalidQuantumGraphFileError( 

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

783 ) 

784 return data 

785 

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

787 """Read a single compressed JSON block. 

788 

789 Parameters 

790 ---------- 

791 address : `Address` 

792 Size and offset of the block. 

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

794 Pydantic model to validate JSON with. 

795 decompressor : `Decompressor` 

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

797 

798 Returns 

799 ------- 

800 model : `pydantic.BaseModel` 

801 Validated model. 

802 """ 

803 compressed_data = self.read_bytes(address) 

804 if compressed_data is None: 

805 return None 

806 json_data = decompressor.decompress(compressed_data) 

807 return model_type.model_validate_json(json_data)