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