Coverage for python/lsst/images/serialization/_io.py: 34%
82 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-08 08:45 +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`` / ``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 "write",
22)
24import importlib
25import importlib.metadata
26from typing import Any, overload
28from lsst.resources import ResourcePathExpression
30from ._backends import backend_for_path
31from ._common import ArchiveReadError, ArchiveTree
33_REGISTRY: dict[str, type[ArchiveTree]] = {}
34"""Map of ``SCHEMA_NAME`` to the registered ``ArchiveTree`` subclass.
36The registry is keyed by name only. Schema-version compatibility is
37enforced when the selected tree's ``model_validate*`` runs, via
38``min_read_version``.
39"""
41_SCHEMA_ENTRY_POINT_GROUP = "lsst.images.schemas"
42"""Entry point group for third-party serialization-model providers."""
44_BUILTIN_SCHEMA_PROVIDERS: dict[str, str] = {
45 "cell_coadd": "lsst.images.cells._coadd:CellCoaddSerializationModel",
46 "cell_psf": "lsst.images.cells._psf:CellPointSpreadFunctionSerializationModel",
47 "coadd_provenance": "lsst.images.cells._provenance:CoaddProvenanceSerializationModel",
48}
49"""Schema providers owned by this package but not imported by ``lsst.images``.
51These duplicate the package's own ``lsst.images.schemas`` entry points so
52source-tree use via ``PYTHONPATH=python`` has the same lazy-import behavior as
53an installed distribution with entry point metadata.
55Schemas whose model classes are imported unconditionally by ``lsst.images`` do
56not need built-in providers or entry points: their ``ArchiveTree`` subclass
57hooks register them before this lazy path is needed.
58"""
61def class_for_schema(schema_name: str) -> type[ArchiveTree] | None:
62 """Return the registered ``ArchiveTree`` subclass for ``schema_name``.
64 If no class is already registered, this attempts schema-specific lazy
65 imports from built-in providers and then from entry points in the
66 ``lsst.images.schemas`` group before returning `None`.
68 Parameters
69 ----------
70 schema_name
71 Schema name (e.g. ``"visit_image"``).
72 """
73 if (cls := _REGISTRY.get(schema_name)) is not None:
74 return cls
75 _load_builtin_schema_provider(schema_name)
76 if (cls := _REGISTRY.get(schema_name)) is not None:
77 return cls
78 _load_schema_entry_points(schema_name)
79 return _REGISTRY.get(schema_name)
82def register_schema_class(cls: type[ArchiveTree]) -> None:
83 """Register ``cls`` under ``cls.SCHEMA_NAME``.
85 No-op when the same class is registered again (re-import during
86 tests). Raises `RuntimeError` when a *different* class is
87 registered under an existing name.
89 Intended to be called from ``ArchiveTree.__pydantic_init_subclass__``;
90 not part of the public API.
91 """
92 key = cls.SCHEMA_NAME
93 existing = _REGISTRY.get(key)
94 if existing is cls: 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 return
96 if existing is not None: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true
97 raise RuntimeError(
98 f"Schema {cls.SCHEMA_NAME!r} is already registered to "
99 f"{existing.__qualname__}; refusing to replace it with "
100 f"{cls.__qualname__}."
101 )
102 _REGISTRY[key] = cls
105def _load_builtin_schema_provider(schema_name: str) -> None:
106 """Import a package-local provider for ``schema_name``, if one exists."""
107 provider = _BUILTIN_SCHEMA_PROVIDERS.get(schema_name)
108 if provider is None:
109 return
110 try:
111 obj = _load_provider_object(provider)
112 except Exception as err:
113 raise ArchiveReadError(
114 f"Could not load built-in schema provider {provider!r} for schema {schema_name!r}: {err}"
115 ) from err
116 _register_provider_object(obj)
117 if schema_name not in _REGISTRY:
118 raise ArchiveReadError(
119 f"Built-in schema provider {provider!r} did not register schema {schema_name!r}."
120 )
123def _load_schema_entry_points(schema_name: str) -> None:
124 """Load entry points named ``schema_name`` from ``lsst.images.schemas``."""
125 loaded: list[str] = []
126 for entry_point in importlib.metadata.entry_points(
127 group=_SCHEMA_ENTRY_POINT_GROUP,
128 name=schema_name,
129 ):
130 loaded.append(entry_point.value)
131 try:
132 obj = entry_point.load()
133 except Exception as err:
134 raise ArchiveReadError(
135 f"Could not load schema provider entry point {entry_point.value!r} "
136 f"for schema {schema_name!r}: {err}"
137 ) from err
138 _register_provider_object(obj)
139 if loaded and schema_name not in _REGISTRY:
140 raise ArchiveReadError(
141 f"Schema provider entry point(s) for {schema_name!r} did not register that schema: {loaded}."
142 )
145def _load_provider_object(provider: str) -> object:
146 """Load ``module[:attribute[.nested]]`` provider specifications."""
147 module_name, _, attr_path = provider.partition(":")
148 obj: object = importlib.import_module(module_name)
149 if attr_path:
150 for attr in attr_path.split("."):
151 obj = getattr(obj, attr)
152 return obj
155def _register_provider_object(obj: object) -> None:
156 """Register ``obj`` if a provider returned an ``ArchiveTree`` subclass."""
157 if isinstance(obj, type) and issubclass(obj, ArchiveTree):
158 register_schema_class(obj)
161def parameterize_tree(
162 tree_cls: type[ArchiveTree],
163 pointer_type: type[Any],
164) -> type[ArchiveTree]:
165 """Parameterise ``tree_cls`` over ``pointer_type`` if it is generic.
167 Some `ArchiveTree` subclasses (e.g. ``SumFieldSerializationModel``)
168 take no type parameters; their ``_get_archive_tree_type`` returns
169 the class itself. Match that behaviour here so per-backend
170 ``open_tree`` implementations can call this uniformly.
171 """
172 if not getattr(tree_cls, "__parameters__", ()):
173 return tree_cls
174 return tree_cls[pointer_type] # type: ignore[index]
177def public_type_for_schema(schema_name: str) -> type | None:
178 """Return the in-memory Python class produced when reading an archive
179 whose top-level tree has schema name ``schema_name``.
181 Looks the schema name up in the registry and returns the registered
182 tree's ``PUBLIC_TYPE`` ClassVar (the type its ``deserialize`` produces).
183 Returns `None` when nothing is registered for ``schema_name``.
185 Parameters
186 ----------
187 schema_name
188 Schema name (e.g. ``"visit_image"``).
189 """
190 tree_cls = class_for_schema(schema_name)
191 if tree_cls is None:
192 return None
193 return getattr(tree_cls, "PUBLIC_TYPE", None)
196@overload
197def read[T](path: ResourcePathExpression, cls: type[T], **kwargs: Any) -> T: ...
198@overload
199def read(path: ResourcePathExpression, cls: None = ..., **kwargs: Any) -> Any: ... 199 ↛ exitline 199 didn't return from function 'read' because
200def read(path, cls=None, **kwargs):
201 """Read an archive whose in-memory type is inferred from its schema.
203 Dispatches to the appropriate backend based on ``path``'s extension,
204 resolves the registered in-memory type from the file's schema, and
205 returns the fully deserialized object.
206 Schema-version compatibility is enforced when the model validates the
207 on-disk tree, via ``min_read_version``.
209 This is the convenient way to read a whole object. To read individual
210 components, or to reach the metadata and butler info stored alongside the
211 object, use `open` instead.
213 Parameters
214 ----------
215 path
216 File to read; convertible to `lsst.resources.ResourcePath`.
217 cls
218 Optional expected in-memory type.
219 When given, the file's schema is checked against ``cls`` and the
220 deserialized object is validated with ``isinstance`` (raising
221 `TypeError` otherwise), and the static return type is ``T``.
222 **kwargs
223 Type-specific keyword arguments forwarded to the object's
224 ``deserialize`` (e.g. ``bbox`` for an image subset read).
225 Mis-targeted arguments surface as ``TypeError``.
226 Backend-specific open options (e.g. ``page_size``) are not accepted
227 here; use `open` for those.
229 Returns
230 -------
231 object
232 The deserialized object.
234 Raises
235 ------
236 ValueError
237 Raised by `backend_for_path` if the file extension is not recognized.
238 ArchiveReadError
239 Raised when the file's ``schema_name`` is not registered, or
240 propagated from the model's ``min_read_version`` check on
241 ``model_validate*``.
242 TypeError
243 Raised when ``cls`` is given and the file's schema or the
244 deserialized object is not compatible with it.
245 """
246 # Imported here to break the _io <-> _reader import cycle: _reader imports
247 # class_for_schema / public_type_for_schema from this module at load time.
248 from ._reader import open as open_archive
250 # A subset read (any deserialize kwarg with a value) reads incrementally;
251 # a plain whole-object read may slurp the file up front. This mirrors the
252 # ``partial`` default the per-backend readers used.
253 partial = any(value is not None for value in kwargs.values())
254 with open_archive(path, cls, partial=partial) as reader:
255 return reader.read(**kwargs)
258def write(obj: Any, path: str, **kwargs: Any) -> Any:
259 """Write ``obj`` to ``path``, dispatching by file extension.
261 Forwards ``**kwargs`` to the per-backend ``write`` (e.g.
262 ``compression_options`` for FITS). No registry lookup is performed:
263 the per-backend ``write`` already accepts any object with a
264 ``serialize`` method.
266 Parameters
267 ----------
268 obj
269 Object to write; must implement ``serialize`` like the per-backend
270 write functions expect.
271 path
272 Destination path. The extension selects the backend.
273 **kwargs
274 Forwarded verbatim to the backend's ``write``.
276 Returns
277 -------
278 Any
279 Whatever the per-backend ``write`` returns (the serialised
280 archive tree).
281 """
282 return backend_for_path(path).write(obj, path, **kwargs)