Coverage for python/lsst/daf/butler/dimensions/

1# This file is part of daf_butler.

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.

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/>.

28from __future__ import annotations

30__all__ = ("DimensionRecord", "SerializedDimensionRecord", "SerializedKeyValueDimensionRecord")

32import itertools

33from collections.abc import Callable, Hashable

34from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeAlias, cast

36import pydantic

37from pydantic import BaseModel, Field, StrictBool, StrictFloat, StrictInt, StrictStr, create_model

39import lsst.sphgeom

40from lsst.utils.classes import immutable

42from .._timespan import Timespan

43from ..column_spec import make_tuple_type_adapter

44from ..json import from_json_pydantic, to_json_pydantic

45from ..persistence_context import PersistenceContextVars

46from ._elements import Dimension, DimensionElement

48if TYPE_CHECKING: # Imports needed only for type annotations; may be circular.

49 from ..registry import Registry

50 from ._coordinate import DataCoordinate

51 from ._schema import DimensionElementFields

52 from ._universe import DimensionUniverse

55DataIdKey: TypeAlias = str

56"""Type annotation alias for the keys that can be used to index a

57DataCoordinate.

58"""

60# Pydantic will cast int to str if str is first in the Union.

61DataIdValue: TypeAlias = int | str

62"""Type annotation alias for the values that can be present in a

63DataCoordinate or other data ID.

64"""

66SerializedKeyValueDimensionRecord: TypeAlias = list[Any]

67"""Type annotation alias for the serialized form of DimensionRecord used in

68container serialization (e.g. `DimensionRecordSet`).

69"""

72def _reconstructDimensionRecord(definition: DimensionElement, mapping: dict[str, Any]) -> DimensionRecord:

73 """Unpickle implementation for `DimensionRecord` subclasses.

75 For internal use by `DimensionRecord`.

76 """

77 return definition.RecordClass(**mapping)

80def _subclassDimensionRecord(definition: DimensionElement) -> type[DimensionRecord]:

81 """Create a dynamic subclass of `DimensionRecord` for the given element.

83 For internal use by `DimensionRecord`.

84 """

85 from ._schema import DimensionElementFields

87 fields = DimensionElementFields(definition)

88 slots = list(fields.standard.names)

89 if definition.spatial:

90 slots.append("region")

91 if definition.temporal:

92 slots.append("timespan")

94 key_type_adapter = make_tuple_type_adapter(definition.schema.required)

95 value_type_adapter = make_tuple_type_adapter(

96 itertools.chain(definition.schema.implied, definition.schema.remainder)

97 )

99 d = {

100 "definition": definition,

101 "__slots__": tuple(slots),

102 "fields": fields,

103 "_key_type_adapter": key_type_adapter,

104 "_value_type_adapter": value_type_adapter,

105 }

106 return type(definition.name + ".RecordClass", (DimensionRecord,), d)

107

108

109class SpecificSerializedDimensionRecord(BaseModel, extra="forbid"):

110 """Base model for a specific serialized record content."""

111

112

113_SIMPLE_RECORD_CLASS_CACHE: dict[

114 tuple[DimensionElement, DimensionUniverse], type[SpecificSerializedDimensionRecord]

115] = {}

116

117

118def _createSimpleRecordSubclass(definition: DimensionElement) -> type[SpecificSerializedDimensionRecord]:

119 from ._schema import DimensionElementFields

120

121 # Cache on the definition (which hashes as the name) and the

122 # associated universe.

123 cache_key = (definition, definition.universe)

124 if cache_key in _SIMPLE_RECORD_CLASS_CACHE:

125 return _SIMPLE_RECORD_CLASS_CACHE[cache_key]

126

127 fields = DimensionElementFields(definition)

128 members = {}

129 # Prefer strict typing for external data

130 type_map = {

131 str: StrictStr,

132 float: StrictFloat,

133 bool: StrictBool,

134 int: StrictInt,

135 }

136

137 for field in fields.standard:

138 field_type = field.getPythonType()

139 field_type = type_map.get(field_type, field_type)

140 if field.nullable:

141 field_type = field_type | None # type: ignore

142 members[field.name] = (field_type, ...)

143 if definition.temporal:

144 members["timespan"] = (Timespan | None, ...) # type: ignore

145 if definition.spatial:

146 members["region"] = (str | None, ...) # type: ignore

147

148 # For the new derived class name need to convert to camel case.

149 # so "day_obs" -> "DayObs".

150 derived_name = "".join([part.capitalize() for part in definition.name.split("_")])

151

152 model = create_model(

153 f"SpecificSerializedDimensionRecord{derived_name}",

154 __base__=SpecificSerializedDimensionRecord,

155 **members, # type: ignore

156 )

157

158 _SIMPLE_RECORD_CLASS_CACHE[cache_key] = model

159 return model

160

161

162# While supporting pydantic v1 and v2 keep this outside the model.

163_serialized_dimension_record_schema_extra = {

164 "examples": [

165 {

166 "definition": "detector",

167 "record": {

168 "instrument": "HSC",

169 "id": 72,

170 "full_name": "0_01",

171 "name_in_raft": "01",

172 "raft": "0",

173 "purpose": "SCIENCE",

174 },

175 }

176 ]

177}

178

179

180class SerializedDimensionRecord(BaseModel):

181 """Simplified model for serializing a `DimensionRecord`."""

182

183 definition: str = Field(

184 ...,

185 title="Name of dimension associated with this record.",

186 examples=["exposure"],

187 )

188

189 # Use strict types to prevent casting

191 ...,

192 title="Dimension record keys and values.",

193 examples=[

194 {

195 "definition": "exposure",

196 "record": {

197 "instrument": "LATISS",

198 "exposure": 2021050300044,

199 "obs_id": "AT_O_20210503_00044",

200 },

201 }

202 ],

203 )

204

205 model_config = {

206 "json_schema_extra": _serialized_dimension_record_schema_extra, # type: ignore[typeddict-item]

207 }

208

209 @classmethod

210 def direct(

211 cls,

212 *,

213 definition: str,

214 record: dict[str, Any],

215 ) -> SerializedDimensionRecord:

216 """Construct a `SerializedDimensionRecord` directly without validators.

217

218 Parameters

219 ----------

220 definition : `str`

221 The name of the record.

222 record : `dict`

223 A dictionary representation of the record content.

224

225 Returns

226 -------

227 rec : `SerializedDimensionRecord`

228 A model representing the dimension records.

229

230 Notes

231 -----

232 This differs from the pydantic "construct" method in that the arguments

233 are explicitly what the model requires, and it will recurse through

234 members, constructing them from their corresponding `direct` methods.

235

236 This method should only be called when the inputs are trusted.

237 """

238 # This method requires tuples as values of the mapping, but JSON

239 # readers will read things in as lists. Be kind and transparently

240 # transform to tuples.

241 _recItems = {

242 k: (v if type(v) is not list else Timespan(begin=None, end=None, _nsec=tuple(v))) # type: ignore

243 for k, v in record.items()

244 }

245

246 # Type ignore because the ternary statement seems to confuse mypy

247 # based on conflicting inferred types of v.

248 key = (

249 definition,

250 frozenset(_recItems.items()),

251 )

252 cache = PersistenceContextVars.serializedDimensionRecordMapping.get()

253 if cache is not None and (result := cache.get(key)) is not None:

254 return result

255

256 node = cls.model_construct(definition=definition, record=_recItems) # type: ignore

257

258 if cache is not None:

259 cache[key] = node

260 return node

261

262

263@immutable

264class DimensionRecord:

265 """Base class for the Python representation of database records.

266

267 Parameters

268 ----------

269 **kwargs

270 Field values for this record. Unrecognized keys are ignored. If this

271 is the record for a `Dimension`, its primary key value may be provided

272 with the actual name of the field (e.g. "id" or "name"), the name of

273 the `Dimension`, or both. If this record class has a "timespan"

274 attribute, "datetime_begin" and "datetime_end" keyword arguments may

275 be provided instead of a single "timespan" keyword argument (but are

276 ignored if a "timespan" argument is provided).

277

278 Notes

279 -----

280 `DimensionRecord` subclasses are created dynamically for each

281 `DimensionElement` in a `DimensionUniverse`, and are accessible via the

282 `DimensionElement.RecordClass` attribute. The `DimensionRecord` base class

283 itself is pure abstract, but does not use the `abc` module to indicate this

284 because it does not have overridable methods.

285

286 Record classes have attributes that correspond exactly to the

287 `~DimensionElementFields.standard` fields in the related database table,

288 plus "region" and "timespan" attributes for spatial and/or temporal

289 elements (respectively).

290

291 Instances are usually obtained from a `Registry`, but can be constructed

292 directly from Python as well.

293

294 `DimensionRecord` instances are immutable.

295 """

296

297 # Derived classes are required to define __slots__ as well, and it's those

298 # derived-class slots that other methods on the base class expect to see

299 # when they access self.__slots__.

300 __slots__ = ("dataId",)

301

302 _serializedType: ClassVar[type[BaseModel]] = SerializedDimensionRecord

303

304 _key_type_adapter: ClassVar[pydantic.TypeAdapter[tuple[Any, ...]]]

305 _value_type_adapter: ClassVar[pydantic.TypeAdapter[tuple[Any, ...]]]

306

307 def __init__(self, **kwargs: Any):

308 # Accept either the dimension name or the actual name of its primary

309 # key field; ensure both are present in the dict for convenience below.

310 if isinstance(self.definition, Dimension):

311 v = kwargs.get(self.definition.primaryKey.name)

312 if v is None:

313 v = kwargs.get(self.definition.name)

314 if v is None:

315 raise ValueError(

316 f"No value provided for {self.definition.name}.{self.definition.primaryKey.name}."

317 )

318 kwargs[self.definition.primaryKey.name] = v

319 else:

320 v2 = kwargs.setdefault(self.definition.name, v)

321 if v != v2:

322 raise ValueError(

323 "Multiple inconsistent values for "

324 f"{self.definition.name}.{self.definition.primaryKey.name}: {v!r} != {v2!r}."

325 )

326

327 from ._coordinate import DataCoordinate

328

329 object.__setattr__(

330 self,

331 "dataId",

332 DataCoordinate.from_required_values(

333 self.definition.minimal_group,

334 tuple(kwargs[dimension] for dimension in self.definition.required.names),

335 ),

336 )

337 # Don't need the primary key value aliased to the dimension name

338 # anymore.

339 kwargs.pop(self.definition.name, None)

340

341 for name in self.__slots__:

342 # Note that we remove from kwargs as we go, to make sure there's

343 # nothing left at the end.

344 object.__setattr__(self, name, kwargs.pop(name, None))

345 # Support 'datetime_begin' and 'datetime_end' instead of 'timespan' for

346 # backwards compatibility, but if one is provided both must be.

347 if self.definition.temporal is not None and self.timespan is None and "datetime_begin" in kwargs:

348 object.__setattr__(

349 self,

350 "timespan",

351 Timespan(

352 kwargs.pop("datetime_begin"),

353 kwargs.pop("datetime_end"),

354 ),

355 )

356

357 if kwargs:

358 raise TypeError(f"Invalid fields for {self.definition} dimension record: {set(kwargs.keys())}.")

359

360 def __eq__(self, other: Any) -> bool:

361 if type(other) is not type(self):

362 return False

363 return all(getattr(self, name) == getattr(other, name) for name in self.__slots__)

364

365 def __hash__(self) -> int:

366 return hash(self.dataId.required_values)

367

368 def __str__(self) -> str:

369 lines = [f"{self.definition.name}:"]

370 lines.extend(f" {name}: {getattr(self, name)!r}" for name in self.__slots__)

371 return "\n".join(lines)

372

373 def __repr__(self) -> str:

374 return "{}.RecordClass({})".format(

375 self.definition.name, ", ".join(f"{name}={getattr(self, name)!r}" for name in self.__slots__)

376 )

377

378 def __reduce__(self) -> tuple:

379 mapping = {name: getattr(self, name) for name in self.__slots__}

380 return (_reconstructDimensionRecord, (self.definition, mapping))

381

382 def _repr_html_(self) -> str:

383 """Override the default representation in IPython/Jupyter notebooks.

384

385 This gives a more readable output that understands embedded newlines.

386 """

387 return f"<pre>{self}<pre>"

388

389 def to_simple(self, minimal: bool = False) -> SerializedDimensionRecord:

390 """Convert this class to a simple python type.

391

392 This makes it suitable for serialization.

393

394 Parameters

395 ----------

396 minimal : `bool`, optional

397 Use minimal serialization. Has no effect on for this class.

398

399 Returns

400 -------

401 names : `list`

402 The names of the dimensions.

403 """

404 # The DataId is sufficient if you are willing to do a deferred

405 # query. This may not be overly useful since to reconstruct

406 # a collection of records will require repeated registry queries.

407 # For now do not implement minimal form.

408 key = (id(self.definition), self.dataId)

409 cache = PersistenceContextVars.serializedDimensionRecordMapping.get()

410 if cache is not None and (result := cache.get(key)) is not None:

411 return result

412

413 mapping = {name: getattr(self, name) for name in self.__slots__}

414 for k, v in mapping.items():

415 if isinstance(v, lsst.sphgeom.Region):

416 # YAML serialization specifies the class when it

417 # doesn't have to. This is partly for explicitness

418 # and also history. Here use a different approach.

419 # This code needs to be migrated to sphgeom

420 mapping[k] = v.encode().hex()

421 if isinstance(v, bytes):

422 # We actually can't handle serializing out to bytes for

423 # hash objects, encode it here to a hex string

424 mapping[k] = v.hex()

425 definition = self.definition.to_simple(minimal=minimal)

426 dimRec = SerializedDimensionRecord(definition=definition, record=mapping)

427 if cache is not None:

428 cache[key] = dimRec

429 return dimRec

430

431 @classmethod

432 def from_simple(

433 cls,

434 simple: SerializedDimensionRecord,

435 universe: DimensionUniverse | None = None,

436 registry: Registry | None = None,

437 cacheKey: Hashable | None = None,

438 ) -> DimensionRecord:

439 """Construct a new object from the simplified form.

440

441 This is generally data returned from the `to_simple`

442 method.

443

444 Parameters

445 ----------

446 simple : `SerializedDimensionRecord`

447 Value return from `to_simple`.

448 universe : `DimensionUniverse`

449 The special graph of all known dimensions of which this graph will

450 be a subset. Can be `None` if `Registry` is provided.

451 registry : `lsst.daf.butler.Registry`, optional

452 Registry from which a universe can be extracted. Can be `None`

453 if universe is provided explicitly.

454 cacheKey : `Hashable` or `None`

455 If this is not None, it will be used as a key for any cached

456 reconstruction instead of calculating a value from the serialized

457 format.

458

459 Returns

460 -------

461 record : `DimensionRecord`

462 Newly-constructed object.

463 """

464 if universe is None and registry is None:

465 raise ValueError("One of universe or registry is required to convert names to a DimensionGroup")

466 if universe is None and registry is not None:

467 universe = registry.dimensions

468 if universe is None:

469 # this is for mypy

470 raise ValueError("Unable to determine a usable universe")

471 # Type ignore because the ternary statement seems to confuse mypy

472 # based on conflicting inferred types of v.

473 key = cacheKey or (

474 simple.definition,

475 frozenset(simple.record.items()), # type: ignore

476 )

477 cache = PersistenceContextVars.dimensionRecords.get()

478 if cache is not None and (result := cache.get(key)) is not None:

479 return result

480

481 definition = DimensionElement.from_simple(simple.definition, universe=universe)

482

483 # Create a specialist subclass model with type validation.

484 # This allows us to do simple checks of external data (possibly

485 # sent as JSON) since for now _reconstructDimensionRecord does not

486 # do any validation.

487 record_model_cls = _createSimpleRecordSubclass(definition)

488 record_model = record_model_cls(**simple.record)

489

490 # Region and hash have to be converted to native form; for now assume

491 # that the keys are special. We make the mapping we need to pass to

492 # the DimensionRecord constructor via getattr, because we don't

493 # model_dump re-disassembling things like Timespans that we've already

494 # assembled.

495 mapping = {k: getattr(record_model, k) for k in definition.schema.names}

496

497 if mapping.get("region") is not None:

498 mapping["region"] = lsst.sphgeom.Region.decode(bytes.fromhex(mapping["region"]))

499 if "hash" in mapping:

500 mapping["hash"] = bytes.fromhex(mapping["hash"].decode())

501

502 dimRec = _reconstructDimensionRecord(definition, mapping)

503 if cache is not None:

504 cache[key] = dimRec

505 return dimRec

506

507 to_json = to_json_pydantic

508 from_json: ClassVar[Callable[..., Self]] = cast(Callable[..., Self], classmethod(from_json_pydantic))

509

510 def toDict(self, splitTimespan: bool = False) -> dict[str, Any]:

511 """Return a vanilla `dict` representation of this record.

512

513 Parameters

514 ----------

515 splitTimespan : `bool`, optional

516 If `True` (`False` is default) transform any "timespan" key value

517 from a `Timespan` instance into a pair of regular

518 ("datetime_begin", "datetime_end") fields.

519 """

520 results = {name: getattr(self, name) for name in self.__slots__}

521 if splitTimespan:

522 timespan = results.pop("timespan", None)

523 if timespan is not None:

524 results["datetime_begin"] = timespan.begin

525 results["datetime_end"] = timespan.end

526 return results

527

528 def get(self, name: str) -> Any:

529 """Return a single metadata value associated with this record.

530

531 Parameters

532 ----------

533 name : `str`

534 Key of the metadata value to be retrieved.

535

536 Returns

537 -------

538 value

539 The metadata value.

540

541 Raises

542 ------

543 KeyError

544 If the given name is not a valid key in this dimension record.

545 """

546 if name not in self.__slots__:

547 raise KeyError(f"'{name}' is not a valid record key for dimension '{self.definition.name}'")

548

549 return getattr(self, name)

550

551 def serialize_key_value(self) -> SerializedKeyValueDimensionRecord:

552 """Serialize this record to a `list` that can be sliced into a key

553 (data ID values) / value (everything else) pair.

554

555 Returns

556 -------

557 raw : `list`

558 List of values with JSON-compatible types.

559

560 Notes

561 -----

562 Unlike `to_simple` / `from_simple`, this serialization approach does

563 not encode the `definition` element in the serialized form. This is

564 expected to be serialized separately (e.g. as part of a homogeneous set

565 of dimension records).

566 """

567 key = list(self.dataId.required_values)

568 value = []

569 for name in self.definition.schema.implied.names:

570 value.append(getattr(self, name))

571 for name in self.definition.schema.remainder.names:

572 value.append(getattr(self, name))

573 return key + self._value_type_adapter.dump_python(tuple(value), mode="json")

574

575 @classmethod

576 def deserialize_key(

577 cls, raw: SerializedKeyValueDimensionRecord

578 ) -> tuple[tuple[DataIdValue, ...], SerializedKeyValueDimensionRecord]:

579 """Deserialize just the key slice of the raw `list` serializeation of a

580 dimension record.

581

582 Parameters

583 ----------

584 raw : `list`

585 Serialized list with JSON-compatible types, as returned by

586 `serialize_key_value`.

587

588 Returns

589 -------

590 key : `tuple`

591 Validated tuple of required data ID values that uniquely identify

592 this record, extracted from the head of ``raw``.

593 raw_value : `list`

594 Remaining unvalidated fields.

595 """

596 n = len(cls.definition.minimal_group.required)

597 return cls._key_type_adapter.validate_python(raw[:n]), raw[n:]

598

599 @classmethod

600 def deserialize_value(

601 cls, key: tuple[DataIdValue, ...], raw_value: SerializedKeyValueDimensionRecord

602 ) -> DimensionRecord:

603 """Deserialize the value slice of the raw `list` form of serialized

604 dimension record.

605

606 Parameters

607 ----------

608 key : `tuple`

609 Validated tuple of required data ID values that uniquely identify

610 this record, as returned by `deserialize_key`.

611 raw_value : `list`

612 Serialized list with JSON-compatible types, with just the non-key

613 items, as returned by `deserialize_key`.

614

615 Returns

616 -------

617 record : `DimensionRecord`

618 A fully-validated `DimensionRecord` with this subclass.

619 """

620 from ._coordinate import DataCoordinate

621

622 result = object.__new__(cls) # bypass the usual __init__

623 result.dataId = DataCoordinate.from_required_values(cls.definition.minimal_group, key)

624 value = cls._value_type_adapter.validate_python(raw_value)

625 for name, val in zip(cls.definition.schema.names, key + value):

626 setattr(result, name, val)

627 return result

628

629 # DimensionRecord subclasses are dynamically created, so static type

630 # checkers can't know about them or their attributes. To avoid having to

631 # put "type: ignore", everywhere, add a dummy __getattr__ that tells type

632 # checkers not to worry about missing attributes.

633 def __getattr__(self, name: str) -> Any:

634 raise AttributeError(name)

635

636 # Class attributes below are shadowed by instance attributes, and are

637 # present just to hold the docstrings for those instance attributes.

638

639 dataId: DataCoordinate

640 """A dict-like identifier for this record's primary keys

641 (`DataCoordinate`).

642 """

643

644 definition: ClassVar[DimensionElement]

645 """The `DimensionElement` whose records this class represents

646 (`DimensionElement`).

647 """

648

649 fields: ClassVar[DimensionElementFields]

650 """A categorized view of the fields in this class

651 (`DimensionElementFields`).

652 """

Coverage for python / lsst / daf / butler / dimensions / _records.py: 24%

200 statements