Coverage for python/lsst/images/serialization/_io.py: 63%
88 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-25 09:03 +0000
1# This file is part of lsst-images.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
11"""Generic ``read_archive`` / ``write_archive`` dispatchers and the
12schema-name registry.
13"""
15from __future__ import annotations
17__all__ = (
18 "class_for_schema",
19 "parameterize_tree",
20 "public_type_for_schema",
21 "read_archive",
22 "register_schema_class",
23 "tree_class_for_info",
24 "write_archive",
25)
27import importlib
28import importlib.metadata
29from pathlib import Path
30from typing import IO, TYPE_CHECKING, Any, overload
32from lsst.resources import ResourcePathExpression
34from ._backends import backend_for_path
35from ._common import ArchiveReadError, ArchiveTree
37if TYPE_CHECKING:
38 from ._input_archive import ArchiveInfo
40_REGISTRY: dict[str, type[ArchiveTree]] = {}
41"""Map of ``SCHEMA_NAME`` to the registered ``ArchiveTree`` subclass.
43The registry is keyed by name only. Schema-version compatibility is
44enforced when the selected tree's ``model_validate*`` runs, via
45``min_read_version``.
46"""
48_SCHEMA_ENTRY_POINT_GROUP = "lsst.images.schemas"
49"""Entry point group for third-party serialization-model providers."""
51_BUILTIN_SCHEMA_PROVIDERS: dict[str, str] = {
52 "cell_aperture_correction_map": (
53 "lsst.images.cells._aperture_corrections:CellApertureCorrectionMapSerializationModel"
54 ),
55 "cell_coadd": "lsst.images.cells._coadd:CellCoaddSerializationModel",
56 "cell_psf": "lsst.images.cells._psf:CellPointSpreadFunctionSerializationModel",
57 "coadd_provenance": "lsst.images.cells._provenance:CoaddProvenanceSerializationModel",
58}
59"""Schema providers owned by this package but not imported by ``lsst.images``.
61These duplicate the package's own ``lsst.images.schemas`` entry points so
62source-tree use via ``PYTHONPATH=python`` has the same lazy-import behavior as
63an installed distribution with entry point metadata.
65Schemas whose model classes are imported unconditionally by ``lsst.images`` do
66not need built-in providers or entry points: their ``ArchiveTree`` subclass
67hooks register them before this lazy path is needed.
68"""
71def class_for_schema(schema_name: str) -> type[ArchiveTree] | None:
72 """Return the registered ``ArchiveTree`` subclass for ``schema_name``.
74 If no class is already registered, this attempts schema-specific lazy
75 imports from built-in providers and then from entry points in the
76 ``lsst.images.schemas`` group before returning `None`.
78 Parameters
79 ----------
80 schema_name
81 Schema name (e.g. ``"visit_image"``).
82 """
83 if (cls := _REGISTRY.get(schema_name)) is not None:
84 return cls
85 _load_builtin_schema_provider(schema_name)
86 if (cls := _REGISTRY.get(schema_name)) is not None:
87 return cls
88 _load_schema_entry_points(schema_name)
89 return _REGISTRY.get(schema_name)
92def register_schema_class(cls: type[ArchiveTree]) -> None:
93 """Register ``cls`` under ``cls.SCHEMA_NAME``.
95 Parameters
96 ----------
97 cls
98 The class to register.
100 Notes
101 -----
102 No-op when the same class is registered again (re-import during
103 tests). Raises `RuntimeError` when a *different* class is
104 registered under an existing name.
106 Intended to be called from ``ArchiveTree.__pydantic_init_subclass__``;
107 not part of the public API.
108 """
109 key = cls.SCHEMA_NAME
110 existing = _REGISTRY.get(key)
111 if existing is cls:
112 return
113 if existing is not None:
114 raise RuntimeError(
115 f"Schema {cls.SCHEMA_NAME!r} is already registered to "
116 f"{existing.__qualname__}; refusing to replace it with "
117 f"{cls.__qualname__}."
118 )
119 _REGISTRY[key] = cls
122def _load_builtin_schema_provider(schema_name: str) -> None:
123 """Import a package-local provider for ``schema_name``, if one exists."""
124 provider = _BUILTIN_SCHEMA_PROVIDERS.get(schema_name)
125 if provider is None:
126 return
127 try:
128 obj = _load_provider_object(provider)
129 except Exception as err:
130 raise ArchiveReadError(
131 f"Could not load built-in schema provider {provider!r} for schema {schema_name!r}: {err}"
132 ) from err
133 _register_provider_object(obj)
134 if schema_name not in _REGISTRY: 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true
135 raise ArchiveReadError(
136 f"Built-in schema provider {provider!r} did not register schema {schema_name!r}."
137 )
140def _load_schema_entry_points(schema_name: str) -> None:
141 """Load entry points named ``schema_name`` from ``lsst.images.schemas``."""
142 loaded: list[str] = []
143 for entry_point in importlib.metadata.entry_points(
144 group=_SCHEMA_ENTRY_POINT_GROUP,
145 name=schema_name,
146 ):
147 loaded.append(entry_point.value)
148 try:
149 obj = entry_point.load()
150 except Exception as err:
151 raise ArchiveReadError(
152 f"Could not load schema provider entry point {entry_point.value!r} "
153 f"for schema {schema_name!r}: {err}"
154 ) from err
155 _register_provider_object(obj)
156 if loaded and schema_name not in _REGISTRY: 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true
157 raise ArchiveReadError(
158 f"Schema provider entry point(s) for {schema_name!r} did not register that schema: {loaded}."
159 )
162def _load_provider_object(provider: str) -> object:
163 """Load ``module[:attribute[.nested]]`` provider specifications."""
164 module_name, _, attr_path = provider.partition(":")
165 obj: object = importlib.import_module(module_name)
166 if attr_path: 166 ↛ 169line 166 didn't jump to line 169 because the condition on line 166 was always true
167 for attr in attr_path.split("."):
168 obj = getattr(obj, attr)
169 return obj
172def _register_provider_object(obj: object) -> None:
173 """Register ``obj`` if a provider returned an ``ArchiveTree`` subclass."""
174 if isinstance(obj, type) and issubclass(obj, ArchiveTree): 174 ↛ exitline 174 didn't return from function '_register_provider_object' because the condition on line 174 was always true
175 register_schema_class(obj)
178def tree_class_for_info(info: ArchiveInfo, path: ResourcePathExpression | IO[bytes]) -> type[ArchiveTree]:
179 """Return the registered `ArchiveTree` subclass for ``info``'s schema.
181 Parameters
182 ----------
183 info
184 Basic archive info whose ``schema_name`` selects the tree class.
185 path
186 Path or stream being opened, used only for the error message.
188 Raises
189 ------
190 ArchiveReadError
191 If no class is registered for the schema.
192 """
193 tree_cls = class_for_schema(info.schema_name)
194 if tree_cls is None:
195 raise ArchiveReadError(f"No registered schema {info.schema_name!r}; cannot open {path!r}.")
196 return tree_cls
199def parameterize_tree(
200 tree_cls: type[ArchiveTree],
201 pointer_type: type[Any],
202) -> type[ArchiveTree]:
203 """Parameterise ``tree_cls`` over ``pointer_type`` if it is generic.
205 Some `ArchiveTree` subclasses (e.g. ``SumFieldSerializationModel``)
206 take no type parameters; their ``_get_archive_tree_type`` returns
207 the class itself. Match that behaviour here so per-backend
208 ``open_tree`` implementations can call this uniformly.
210 Parameters
211 ----------
212 tree_cls
213 Archive tree class to parameterise.
214 pointer_type
215 Pointer type to parameterise ``tree_cls`` over when it is generic.
216 """
217 if not getattr(tree_cls, "__parameters__", ()):
218 return tree_cls
219 return tree_cls[pointer_type] # type: ignore[index]
222def public_type_for_schema(schema_name: str) -> type | None:
223 """Return the in-memory Python class produced when reading an archive
224 whose top-level tree has schema name ``schema_name``.
226 Looks the schema name up in the registry and returns the registered
227 tree's ``PUBLIC_TYPE`` ClassVar (the type its ``deserialize`` produces).
228 Returns `None` when nothing is registered for ``schema_name``.
230 Parameters
231 ----------
232 schema_name
233 Schema name (e.g. ``"visit_image"``).
234 """
235 tree_cls = class_for_schema(schema_name)
236 if tree_cls is None:
237 return None
238 return getattr(tree_cls, "PUBLIC_TYPE", None)
241@overload
242def read_archive[T](
243 path: ResourcePathExpression | IO[bytes], cls: type[T], *, format: str | None = ..., **kwargs: Any
244) -> T: ...
245@overload
246def read_archive(
247 path: ResourcePathExpression | IO[bytes], cls: None = ..., *, format: str | None = ..., **kwargs: Any
248) -> Any: ...
249def read_archive(
250 path: ResourcePathExpression | IO[bytes],
251 cls: type[Any] | None = None,
252 *,
253 format: str | None = None,
254 **kwargs: Any,
255) -> Any:
256 """Read an archive whose in-memory type is inferred from its schema.
258 Dispatches to the appropriate backend based on ``path``'s extension (or,
259 for stream input, its leading bytes), resolves the registered in-memory
260 type from the file's schema, and returns the fully deserialized object.
261 Schema-version compatibility is enforced when the model validates the
262 on-disk tree, via ``min_read_version``.
263 A path with a ``.gz``/``.zst`` compression suffix is decompressed
264 transparently; stream input must already be decompressed.
266 This is the convenient way to read a whole object. To read individual
267 components, or to reach the metadata and butler info stored alongside the
268 object, use `open_archive` instead.
270 Parameters
271 ----------
272 path
273 File to read; convertible to `lsst.resources.ResourcePath`, or a
274 seekable binary stream containing the file's content (e.g.
275 ``io.BytesIO(data)`` for in-memory bytes).
276 cls
277 Optional expected in-memory type.
278 When given, the file's schema is checked against ``cls`` and the
279 deserialized object is validated with ``isinstance`` (raising
280 `TypeError` otherwise), and the static return type is ``T``.
281 format
282 Optional backend name (``"fits"``, ``"ndf"``, or ``"json"``)
283 forcing the backend, instead of dispatching on the path's extension
284 or the stream's leading bytes.
285 **kwargs
286 Type-specific keyword arguments forwarded to the object's
287 ``deserialize`` (e.g. ``bbox`` for an image subset read).
288 Mis-targeted arguments surface as ``TypeError``.
289 Backend-specific open options (e.g. ``page_size``) are not accepted
290 here; use `open_archive` for those.
292 Returns
293 -------
294 object
295 The deserialized object.
297 Raises
298 ------
299 ValueError
300 Raised by `backend_for_path` if the file extension is not
301 recognized, by `backend_for_stream` if a stream's leading bytes are
302 not recognized, or by `backend_for_name` if ``format`` is not a
303 known backend name.
304 ArchiveReadError
305 Raised when the file's ``schema_name`` is not registered, or
306 propagated from the model's ``min_read_version`` check on
307 ``model_validate*``.
308 TypeError
309 Raised when ``cls`` is given and the file's schema or the
310 deserialized object is not compatible with it.
311 """
312 # Imported here to break the _io <-> _reader import cycle: _reader imports
313 # class_for_schema / public_type_for_schema from this module at load time.
314 from ._reader import open_archive
316 # A subset read (any deserialize kwarg with a value) reads incrementally;
317 # a plain whole-object read may slurp the file up front. This mirrors the
318 # ``partial`` default the per-backend readers used.
319 partial = any(value is not None for value in kwargs.values())
320 with open_archive(path, cls, format=format, partial=partial) as reader:
321 return reader.read(**kwargs)
324def write_archive(obj: Any, path: Path | str, **kwargs: Any) -> Any:
325 """Write ``obj`` to ``path``, dispatching by file extension.
327 Forwards ``**kwargs`` to the per-backend ``write`` (e.g.
328 ``compression_options`` for FITS). No registry lookup is performed:
329 the per-backend ``write`` already accepts any object with a
330 ``serialize`` method.
332 Parameters
333 ----------
334 obj
335 Object to write; must implement ``serialize`` like the per-backend
336 write functions expect.
337 path
338 Destination path. The extension selects the backend.
339 **kwargs
340 Forwarded verbatim to the backend's ``write``.
342 Returns
343 -------
344 Any
345 Whatever the per-backend ``write`` returns (the serialised
346 archive tree).
347 """
348 return backend_for_path(path).write(obj, path, **kwargs)