Coverage for python/lsst/scarlet/lite/io/hierarchical_blend.py: 56%

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4from typing import Any 

5 

6import numpy as np 

7from numpy.typing import DTypeLike 

8 

9from .blend_base import ScarletBlendBaseData 

10from .migration import PRE_SCHEMA, MigrationRegistry, migration 

11from .utils import PersistenceError, decode_metadata, encode_metadata 

12 

13__all__ = ["HierarchicalBlendData"] 

14 

15CURRENT_SCHEMA = "1.0.0" 

16BLEND_TYPE = "hierarchical" 

17MigrationRegistry.set_current(BLEND_TYPE, CURRENT_SCHEMA) 

18 

19 

20@dataclass(kw_only=True) 

21class HierarchicalBlendData(ScarletBlendBaseData): 

22 """Data for a hierarchical blend. 

23 

24 Attributes 

25 ---------- 

26 blend_type : 

27 The type of blend being stored 

28 children : 

29 Map from blend IDs to child blends. 

30 version : 

31 The schema version of the HierarchicalBlendData. 

32 """ 

33 

34 blend_type: str = BLEND_TYPE 

35 children: dict[int, ScarletBlendBaseData] 

36 version: str = CURRENT_SCHEMA 

37 

38 def as_dict(self) -> dict: 

39 """Return the object encoded into a dict for JSON serialization 

40 

41 Returns 

42 ------- 

43 result : 

44 The object encoded as a JSON compatible dict 

45 """ 

46 result: dict[str, Any] = { 

47 "blend_type": self.blend_type, 

48 "children": {bid: child.as_dict() for bid, child in self.children.items()}, 

49 "version": self.version, 

50 } 

51 if self.metadata is not None: 

52 result["metadata"] = encode_metadata(self.metadata) 

53 return result 

54 

55 @classmethod 

56 def from_dict(cls, data: dict, dtype: DTypeLike = np.float32) -> HierarchicalBlendData: 

57 """Reconstruct `HierarchicalBlendData` from JSON compatible dict. 

58 

59 Parameters 

60 ---------- 

61 data : 

62 Dictionary representation of the object 

63 dtype : 

64 Datatype of the resulting model. 

65 

66 Returns 

67 ------- 

68 result : 

69 The reconstructed object 

70 """ 

71 data = MigrationRegistry.migrate(BLEND_TYPE, data) 

72 children: dict[int, ScarletBlendBaseData] = {} 

73 for blend_id, child in data["children"].items(): 

74 try: 

75 children[int(blend_id)] = ScarletBlendBaseData.from_dict(child, dtype=dtype) 

76 except KeyError: 

77 raise PersistenceError(f"Unknown blend type: {child['blend_type']} for blend ID: {blend_id}") 

78 

79 metadata = decode_metadata(data.get("metadata", None)) 

80 return cls(children=children, metadata=metadata) 

81 

82 

83HierarchicalBlendData.register() 

84# Register the legacy blend_type in the blend registry 

85ScarletBlendBaseData.blend_registry["hierarchical_blend"] = HierarchicalBlendData 

86 

87 

88@migration(BLEND_TYPE, PRE_SCHEMA) 

89def _to_1_0_0(data: dict) -> dict: 

90 """Migrate a pre-schema hierarchical blend to schema version 1.0.0 

91 

92 There were no changes to this data model in v1.0.0 but we need 

93 to provide a way to migrate pre-schema data. 

94 

95 Parameters 

96 ---------- 

97 data : 

98 The data to migrate. 

99 

100 Returns 

101 ------- 

102 result : 

103 The migrated data. 

104 """ 

105 data["version"] = "1.0.0" 

106 return data