pa_api.xmlapi.types.utils

  1import json
  2import logging
  3import typing
  4from datetime import datetime
  5from types import new_class
  6from typing import Annotated, Any, Optional, TypeVar
  7
  8from pydantic import BaseModel, TypeAdapter
  9from pydantic.functional_validators import (
 10    BeforeValidator,
 11    PlainValidator,
 12)
 13from typing_extensions import Self, TypedDict
 14
 15from pa_api.utils import first
 16from pa_api.xmlapi.utils import el2dict
 17
 18TIME_FORMAT = "%H:%M:%S"
 19DATETIME_FORMAT = f"%Y/%m/%d {TIME_FORMAT}"
 20DATETIME_MS_FORMAT = f"{DATETIME_FORMAT}.%f"
 21NoneType: type = type(None)
 22
 23
 24class XMLBaseModel(BaseModel):
 25    @classmethod
 26    def from_xml(cls, xml) -> Self:
 27        data = first(el2dict(xml).values())
 28        return cls.model_validate(data)
 29
 30
 31def parse_datetime(d):
 32    try:
 33        if d is None or d in ("none", "Unknown"):
 34            return None
 35        try:
 36            return datetime.strptime(d, DATETIME_FORMAT)
 37        except Exception:
 38            return datetime.strptime(d, DATETIME_MS_FORMAT)
 39    except Exception as e:
 40        logging.debug(e)
 41        logging.debug(f"Failed to parse {d} as datetime")
 42        # print(d, type(d))
 43        raise
 44    return d
 45
 46
 47def parse_time(d):
 48    return datetime.strptime(d, TIME_FORMAT).time()
 49
 50
 51# https://docs.pydantic.dev/latest/concepts/types/#custom-types
 52# JobProgress = TypeAliasType('JobProgress', PlainValidator(parse_progress))
 53# Datetime = TypeAliasType(
 54#     "Datetime", Annotated[datetime, PlainValidator(parse_datetime)]
 55# )
 56Datetime = Annotated[datetime, PlainValidator(parse_datetime)]
 57
 58
 59def single_xpath(xml, xpath, parser=None, default=None):
 60    try:
 61        res = xml.xpath(xpath)
 62        res = first(res, None)
 63    except Exception:
 64        return default
 65    if res is None:
 66        return default
 67    if not isinstance(res, str):
 68        res = res.text
 69    if parser:
 70        res = parser(res)
 71    return res
 72
 73
 74pd = parse_datetime
 75sx = single_xpath
 76
 77
 78def mksx(xml):
 79    def single_xpath(xpath, parser=None, default=None):
 80        res = sx(xml, xpath, parser=parser, default=default)
 81        logging.debug(res)
 82        return res
 83
 84    return single_xpath
 85
 86
 87def ensure_list(v: Any) -> typing.List[Any]:
 88    if v is None:
 89        return []
 90    if isinstance(v, dict) and len(v) == 1 and "member" in v:
 91        return ensure_list(v["member"])
 92    if isinstance(v, list):
 93        return v
 94    return [v]
 95
 96
 97# https://docs.pydantic.dev/latest/concepts/types/#generics
 98Element = TypeVar("Element", bound=Any)
 99# Similar to typing.List, but ensure to always return a list
100List = Annotated[typing.List[Element], BeforeValidator(ensure_list)]
101
102
103def xml_text(v: Any):
104    if isinstance(v, dict) and "#text" in v:
105        return v["#text"]
106    return v
107
108
109def ensure_str(v: Any) -> str:
110    if v is None:
111        return ""
112    if isinstance(v, dict):
113        return v["#text"]
114    return v
115
116
117def validate_ip(v: Any) -> str:
118    if v is None:
119        return ""
120    if isinstance(v, dict):
121        return v["@name"]
122    return v
123
124
125String = Annotated[str, BeforeValidator(ensure_str)]
126Bool = Annotated[bool, BeforeValidator(xml_text)]
127Ip = Annotated[str, BeforeValidator(validate_ip)]
128
129
130def _xml2schema(values: list):
131    types: typing.List[Any] = [
132        t.__name__ for t in {type(v) for v in values if not isinstance(v, (dict, list))}
133    ]
134    list_values = [v for v in values if isinstance(v, list)]
135    if list_values:
136        types.append(_xml2schema([e for sublist in list_values for e in sublist]))
137    dict_values = [v for v in values if isinstance(v, dict)]
138    if dict_values:
139        all_keys = {k for d in dict_values for k in d}
140        dict_schema = {
141            k: _xml2schema([d.get(k) for d in dict_values]) for k in all_keys
142        }
143        types.append(dict_schema)
144    if not types:
145        raise Exception("NO TYPE")
146    if len(types) == 1:
147        return types[0]
148    return types
149
150
151def _clean_key(k):
152    return k.replace("@", "").replace("#", "")
153
154
155def _slug(k):
156    return _clean_key(k).replace("-", "_")
157
158
159def _keyastypename(k):
160    return "".join(p.title() for p in _clean_key(k).split("-"))
161
162
163def _schematype(values: list, name: Optional[str] = None) -> type:
164    if not name:
165        name = "Dict"
166    optional = None in values
167    values = [v for v in values if v is not None]
168    schema_types: typing.List[Any] = list(
169        {type(v) for v in values if not isinstance(v, (dict, list))}
170    )
171    list_values = [v for v in values if isinstance(v, list)]
172    if list_values:
173        t = _schematype([e for sublist in list_values for e in sublist], name=name)
174        schema_types.append(t)
175    dict_values = [v for v in values if isinstance(v, dict)]
176    if dict_values:
177        all_keys = {k for d in dict_values for k in d}
178        annotations = {
179            _slug(k): _schematype(
180                [d.get(k) for d in dict_values], name=_keyastypename(k)
181            )
182            for k in all_keys
183        }
184        t = new_class(
185            name,
186            (TypedDict,),
187            None,
188            lambda ns: ns.update({"__annotations__": annotations}),
189        )
190        schema_types.append(t)
191    if not schema_types:
192        return NoneType
193        # raise Exception("NO TYPE")
194
195    final_type = (
196        schema_types[0] if len(schema_types) == 1 else typing.Union[tuple(schema_types)]
197    )
198    if optional:
199        final_type = Optional[final_type]
200    return final_type
201
202
203def xml2schema(xml):
204    """
205    Similar to schematype function:
206    The result is a recursive schema that can be dumped with json.
207    """
208    data = el2dict(xml)
209    return _xml2schema([data])
210
211
212def schematype(data, name: Optional[str] = None):
213    """
214    Recursively parse the data to infer the schema.
215    The schema is returned as a type.
216
217    We can dump the json schema using pydantic:
218    TypeAdapter(schematype({"test": 5})).json_schema()
219    """
220    return _schematype([data], name=name)
221
222
223def xml2schematype(xml):
224    """
225    Same to schematype function but takes an xml Element as parameter
226    """
227    name, data = first(el2dict(xml).items())
228    return schematype(data, name)
229
230
231def jsonschema(data, name=None, indent=4) -> str:
232    ta = TypeAdapter(schematype(data, name))
233    return json.dumps(ta.json_schema(), indent=indent)
234
235
236def xml2jsonschema(data, indent=4) -> str:
237    ta = TypeAdapter(xml2schematype(data))
238    return json.dumps(ta.json_schema(), indent=indent)