1# =====================================================
2# Imports
3# =====================================================
4from typing import Annotated
5from xml.etree import cElementTree as et
6
7from pydantic import Field, field_validator, ValidationInfo
8
9import mt_metadata.transfer_functions.io.emtfxml.metadata.helpers as helpers
10from mt_metadata.base import MetadataBase
11from mt_metadata.common import Comment
12from mt_metadata.common.enumerations import ElectrodeLocationEnum
13
14
15# =====================================================
16
17
18class Electrode(MetadataBase):
19 location: Annotated[
20 ElectrodeLocationEnum,
21 Field(
22 default="",
23 description="Direction of electrode",
24 alias=None,
25 json_schema_extra={
26 "units": None,
27 "required": True,
28 "examples": ["N", "S", "E", "W"],
29 },
30 ),
31 ]
32
33 number: Annotated[
34 str,
35 Field(
36 default="0",
37 description="Electrode ID number",
38 alias=None,
39 json_schema_extra={
40 "units": None,
41 "required": True,
42 "examples": ["1a"],
43 },
44 ),
45 ]
46
47 comments: Annotated[
48 Comment,
49 Field(
50 default_factory=lambda: Comment(), # type: ignore[return-value]
51 description="comments on the electrode",
52 alias=None,
53 json_schema_extra={
54 "units": None,
55 "required": True,
56 "examples": ["Ag-AgCl porous pot"],
57 },
58 ),
59 ]
60
61 @field_validator("comments", mode="before")
62 @classmethod
63 def validate_comments(cls, value, info: ValidationInfo) -> Comment:
64 if isinstance(value, str):
65 return Comment(value=value) # type: ignore[return-value]
66 return value
67
68 def to_xml(self, string: bool = False, required: bool = False) -> str | et.Element:
69 """ """
70
71 root = et.Element(
72 self.__class__.__name__,
73 {"location": self.location.upper(), "number": self.number},
74 )
75
76 # this might break in the future when to_dict is updated to return a dict
77 # instead of a string, but for now it works.
78 root.text = self.comments.as_string()
79
80 if string:
81 return helpers.element_to_string(root)
82 return root