diff libervia/backend/models/types.py @ 4324:e1fcf4dd9012

core (models/types): add types for `domish.Element`.
author Goffi <goffi@goffi.org>
date Wed, 20 Nov 2024 11:29:36 +0100
parents ffc43219e0b2
children
line wrap: on
line diff
--- a/libervia/backend/models/types.py	Wed Nov 20 11:28:23 2024 +0100
+++ b/libervia/backend/models/types.py	Wed Nov 20 11:29:36 2024 +0100
@@ -22,6 +22,77 @@
 from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import core_schema
 from twisted.words.protocols.jabber.jid import JID
+from twisted.words.xish.domish import Element
+from libervia.backend.tools.xml_tools import parse
+
+
+class _DomishElementType:
+    """Use Twisted's domish.Element in Python type, and serialize it to str
+
+    In Python, both Element and str are accepted, str are converted to Element.
+    Serialization is done using toXml() method, and parsing using xml_tools.parse.
+    """
+
+    @staticmethod
+    def validate_element(value: str) -> Element:
+        return parse(value)
+
+    @staticmethod
+    def serialize_element(element: Element) -> str:
+        return element.toXml()
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.json_or_python_schema(
+            json_schema=core_schema.no_info_after_validator_function(
+                cls.validate_element, core_schema.str_schema()
+            ),
+            python_schema=core_schema.union_schema(
+                [
+                    core_schema.is_instance_schema(Element),
+                    core_schema.no_info_after_validator_function(
+                        cls.validate_element, core_schema.str_schema()
+                    ),
+                ]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls.serialize_element, when_used="json"
+            ),
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        json_schema = handler(schema)
+        json_schema.update(
+            {
+                "format": "xml",
+                "description": "A valid XML element that can be parsed into a "
+                "domish.Element",
+            }
+        )
+        return json_schema
+
+
+class _StrictDomishElementType(_DomishElementType):
+    """Strict version of DomishElementType which only accept Element in Python."""
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.json_or_python_schema(
+            json_schema=core_schema.no_info_after_validator_function(
+                cls.validate_element, core_schema.str_schema()
+            ),
+            python_schema=core_schema.is_instance_schema(Element),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls.serialize_element, when_used="json"
+            ),
+        )
 
 
 class _JIDType:
@@ -95,3 +166,5 @@
 # We annotate the types so static type checker understand them as JID.
 JIDType = Annotated[JID, _JIDType]
 StrictJIDType = Annotated[JID, _StrictJIDType]
+DomishElementType = Annotated[Element, _DomishElementType]
+StrictDomishElementType = Annotated[Element, _StrictDomishElementType]