Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ transfer_functions \ io \ emtfxml \ metadata \ site_layout.py: 97%

1# ===================================================== 

2# Imports 

3# ===================================================== 

4from typing import Annotated 

5from xml.etree import ElementTree as et 

6 

7from loguru import logger 

8from pydantic import computed_field, Field, field_validator 

9 

10from mt_metadata.base import MetadataBase 

11from mt_metadata.base.helpers import element_to_string 

12 

13from . import Electric, Magnetic 

14 

15 

16# ===================================================== 

17class SiteLayout(MetadataBase): 

18 input_channels: Annotated[ 

19 list[Electric | Magnetic | str], 

20 Field( 

21 default_factory=list, 

22 description="list of input channels for transfer function estimation", 

23 alias=None, 

24 json_schema_extra={ 

25 "units": None, 

26 "required": True, 

27 "examples": ["[Magnetic(hx), Magnetic(hy)]"], 

28 }, 

29 ), 

30 ] 

31 

32 output_channels: Annotated[ 

33 list[Electric | Magnetic | str], 

34 Field( 

35 default_factory=list, 

36 description="list of output channels for transfer function estimation", 

37 alias=None, 

38 json_schema_extra={ 

39 "units": None, 

40 "required": True, 

41 "examples": ["[Electric(ex), Electric(ey), Magnetic(hz)]"], 

42 }, 

43 ), 

44 ] 

45 

46 @computed_field 

47 @property 

48 def input_channel_names(self) -> list[str]: 

49 """ 

50 Returns a list of input channel names. 

51 """ 

52 return [ch.name.lower() for ch in self.input_channels] 

53 

54 @computed_field 

55 @property 

56 def output_channel_names(self) -> list[str]: 

57 """ 

58 Returns a list of output channel names. 

59 """ 

60 return [ch.name.lower() for ch in self.output_channels] 

61 

62 @field_validator("input_channels", "output_channels", mode="before") 

63 @classmethod 

64 def validate_channels( 

65 cls, value: list[Electric | Magnetic | str] 

66 ) -> list[Electric | Magnetic]: 

67 channels = [] 

68 if not isinstance(value, list): 

69 value = [value] 

70 

71 for item in value: 

72 if isinstance(item, (Magnetic, Electric)): 

73 channels.append(item) 

74 elif isinstance(item, dict): 

75 try: 

76 # Assume the dict has a single key for channel type 

77 ch_type = list(item.keys())[0] 

78 except IndexError: 

79 msg = "Channel dict must have a single key for channel type" 

80 logger.error(msg) 

81 raise ValueError(msg) 

82 ch_type = list(item.keys())[0] 

83 if ch_type in ["magnetic"]: 

84 ch = Magnetic() # type: ignore 

85 elif ch_type in ["electric"]: 

86 ch = Electric() # type: ignore 

87 else: 

88 msg = f"Channel type {ch_type} not supported" 

89 logger.error(msg) 

90 raise ValueError(msg) 

91 ch.from_dict(item) 

92 channels.append(ch) 

93 elif isinstance(item, str): 

94 if item.lower().startswith("e"): 

95 ch = Electric(name=item) # type: ignore 

96 elif item.lower().startswith("b") or item.lower().startswith("h"): 

97 ch = Magnetic(name=item) # type: ignore 

98 else: 

99 msg = f"Channel {item} not supported" 

100 logger.error(msg) 

101 raise ValueError(msg) 

102 channels.append(ch) 

103 else: 

104 msg = f"Channel {item} not supported" 

105 logger.error(msg) 

106 raise TypeError(msg) 

107 

108 return channels 

109 

110 def read_dict(self, input_dict: dict) -> None: 

111 """ 

112 read site layout into the proper input/output channels 

113 

114 :param input_dict: input dictionary containing site layout data 

115 :type input_dict: dict 

116 :return: None 

117 :rtype: None 

118 

119 """ 

120 # read input channels 

121 for ch in ["input_channels", "output_channels"]: 

122 ch_list = [] 

123 try: 

124 c_list = input_dict["site_layout"][ch]["magnetic"] 

125 if c_list is None: 

126 continue 

127 if not isinstance(c_list, list): 

128 c_list = [c_list] 

129 ch_list += [{"magnetic": ch_dict} for ch_dict in c_list] 

130 

131 except (KeyError, TypeError): 

132 pass 

133 

134 try: 

135 c_list = input_dict["site_layout"][ch]["electric"] 

136 if c_list is None: 

137 continue 

138 if not isinstance(c_list, list): 

139 c_list = [c_list] 

140 ch_list += [{"electric": ch_dict} for ch_dict in c_list] 

141 except (KeyError, TypeError): 

142 pass 

143 

144 setattr(self, ch, ch_list) 

145 

146 def to_xml(self, string: bool = False, required: bool = True) -> str | et.Element: 

147 """ 

148 Convert the SiteLayout instance to an XML representation. 

149 

150 Parameters 

151 ---------- 

152 string : bool, optional 

153 Whether to return the XML as a string, by default False 

154 required : bool, optional 

155 Whether the XML elements are required, by default True 

156 

157 Returns 

158 ------- 

159 str | et.Element 

160 The XML representation of the SiteLayout instance 

161 """ 

162 

163 root = et.Element(self.__class__.__name__) 

164 

165 section = et.SubElement( 

166 root, "InputChannels", attrib={"ref": "site", "units": "m"} 

167 ) 

168 for ch in self.input_channels: 

169 section.append(ch.to_xml(required=required)) 

170 section = et.SubElement( 

171 root, "OutputChannels", attrib={"ref": "site", "units": "m"} 

172 ) 

173 for ch in self.output_channels: 

174 section.append(ch.to_xml(required=required)) 

175 

176 if string: 

177 return element_to_string(root) 

178 return root