Coverage for python/lsst/images/json/_output_archive.py: 59%
60 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-01 09:12 +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.
12from __future__ import annotations
14__all__ = ("JsonOutputArchive", "write")
16from collections.abc import Callable, Hashable, Iterator, Mapping
17from typing import TYPE_CHECKING, Any
19import astropy.table
20import numpy as np
21import pydantic
23from lsst.resources import ResourcePath, ResourcePathExpression
25from .._transforms import FrameSet
26from ..serialization import (
27 ArchiveTree,
28 ButlerInfo,
29 InlineArrayModel,
30 JsonRef,
31 MetadataValue,
32 NestedOutputArchive,
33 NumberType,
34 OutputArchive,
35 TableColumnModel,
36 TableModel,
37 no_header_updates,
38)
40if TYPE_CHECKING:
41 import astropy.io.fits
44def write(
45 obj: Any,
46 path: ResourcePathExpression | None = None,
47 metadata: dict[str, MetadataValue] | None = None,
48 butler_info: ButlerInfo | None = None,
49) -> ArchiveTree:
50 """Write an object with a ``serialize`` method to a JSON file.
52 Parameters
53 ----------
54 obj
55 Object with a ``serialize`` method to write.
56 path
57 Name of the file to write to (may be a URI). If not provided, a
58 serializable model is returned but not written to disk.
59 metadata
60 Additional metadata to save with the object. This will override any
61 flexible metadata carried by the object itself with the same keys.
62 butler_info
63 Butler information to store in the file.
65 Returns
66 -------
67 `.serialization.ArchiveTree`
68 The serialized representation of the object.
69 """
70 archive = JsonOutputArchive()
71 name = getattr(obj, "_archive_default_name", None)
72 tree = archive.serialize_direct(name, obj.serialize) if name is not None else obj.serialize(archive)
73 if metadata is not None: 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true
74 tree.metadata.update(metadata)
75 if butler_info is not None:
76 tree.butler_info = butler_info
77 archive.finish(tree)
78 if path is not None: 78 ↛ 80line 78 didn't jump to line 80 because the condition on line 78 was always true
79 ResourcePath(path).write(tree.model_dump_json().encode())
80 return tree
83class JsonOutputArchive(OutputArchive[JsonRef]):
84 """An implementation of the `.serialization.OutputArchive` interface that
85 writes to JSON files.
87 This archive type is designed for pure-JSON objects and cases where any
88 images or tables are tiny. It will be *extremely* inefficient for large
89 images or tables, if it works at all.
90 """
92 def __init__(self) -> None:
93 super().__init__()
94 self._pointers: dict[Hashable, JsonRef] = {}
95 self._indirect: list[Any] = []
96 self._frame_sets: list[tuple[FrameSet, JsonRef]] = []
98 def serialize_direct[T: pydantic.BaseModel | None](
99 self, name: str, serializer: Callable[[OutputArchive[JsonRef]], T]
100 ) -> T:
101 nested = NestedOutputArchive[JsonRef](name, self)
102 return serializer(nested)
104 def serialize_pointer[T: ArchiveTree](
105 self, name: str, serializer: Callable[[OutputArchive[JsonRef]], T], key: Hashable
106 ) -> JsonRef:
107 if (pointer := self._pointers.get(key)) is not None:
108 return pointer
109 pointer = JsonRef.model_construct(ref=f"#/indirect/{len(self._indirect)}")
110 self._indirect.append(self.serialize_direct(name, serializer).model_dump())
111 self._pointers[key] = pointer
112 return pointer
114 def serialize_frame_set[T: ArchiveTree](
115 self, name: str, frame_set: FrameSet, serializer: Callable[[OutputArchive], T], key: Hashable
116 ) -> JsonRef:
117 pointer = self.serialize_pointer(name, serializer, key)
118 self._frame_sets.append((frame_set, pointer))
119 return pointer
121 def iter_frame_sets(self) -> Iterator[tuple[FrameSet, JsonRef]]:
122 return iter(self._frame_sets)
124 def add_array(
125 self,
126 array: np.ndarray,
127 *,
128 name: str | None = None,
129 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
130 tile_shape: tuple[int, ...] | None = None,
131 options_name: str | None = None,
132 ) -> InlineArrayModel:
133 return InlineArrayModel(
134 data=array.tolist(),
135 datatype=NumberType.from_numpy(array.dtype),
136 )
138 def add_table(
139 self,
140 table: astropy.table.Table,
141 *,
142 name: str | None = None,
143 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
144 ) -> TableModel:
145 columns = TableColumnModel.from_table(table, inline=True)
146 return TableModel(columns=columns, meta=table.meta)
148 def add_structured_array(
149 self,
150 array: np.ndarray,
151 *,
152 name: str | None = None,
153 units: Mapping[str, astropy.units.Unit] | None = None,
154 descriptions: Mapping[str, str] | None = None,
155 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
156 ) -> TableModel:
157 columns = TableColumnModel.from_record_array(array, inline=True)
158 for c in columns:
159 if units and (unit := units.get(c.name)):
160 c.unit = unit
161 if descriptions and (description := descriptions.get(c.name)):
162 c.description = description
163 return TableModel(columns=columns)
165 def finish[T: ArchiveTree](self, tree: T) -> T:
166 """Finish serialization.
168 Parameters
169 ----------
170 tree
171 Serialized archive tree to write, which is modified in place
172 (the ``indirect`` attribute is overwritten) and then returned.
173 """
174 tree.indirect = self._indirect
175 return tree