Support customizing tagged union tag values

Author: mickvangelderen

pyproject.toml

@@ -45,6 +45,7 @@ orjson = { version = "*", markers = "extra == 'orjson' or extra == 'all'", optio
 plum-dispatch = ">=2.3"
 beartype = ">=0.18.4"
 sqlalchemy = { version = ">2", markers = "extra == 'sqlalchemy' or extra == 'all'", optional = true }
+frozendict = ">2"
 
 [tool.poetry.dev-dependencies]
 pyyaml = "*"

serde/__init__.py

@@ -35,6 +35,7 @@
     coerce,
     DefaultTagging,
     ExternalTagging,
+    ExternalTagging_,
     InternalTagging,
     disabled,
     strict,
@@ -81,6 +82,7 @@
     "SerdeSkip",
     "AdjacentTagging",
     "ExternalTagging",
+    "ExternalTagging_",
     "InternalTagging",
     "Untagged",
     "disabled",

serde/core.py

@@ -11,6 +11,7 @@
 import re
 import casefy
 from dataclasses import dataclass
+from frozendict import frozendict
 
 from beartype.door import is_bearable
 from collections.abc import Mapping, Sequence, Callable
@@ -776,6 +777,13 @@ class Kind(enum.Enum):
     tag: Optional[str] = None
     content: Optional[str] = None
     kind: Kind = Kind.External
+    tags: Optional[frozendict[type[Any], str]] = None
+
+    def __init__(self, tag: Optional[str] = None, content: Optional[str] = None, kind: Kind = Kind.External, tags: Optional[dict[type[Any], str]] = None) -> None:
+        self.tag = tag
+        self.content = content
+        self.kind = kind
+        self.tags = frozendict(tags) if tags else None
 
     def is_external(self) -> bool:
         return self.kind == self.Kind.External
@@ -793,6 +801,11 @@ def is_untagged(self) -> bool:
     def is_taggable(cls, typ: type[Any]) -> bool:
         return dataclasses.is_dataclass(typ)
 
+    def tag_for(self, typ: type[Any]) -> str:
+        if self.tags and typ in self.tags:
+            return self.tags[typ]
+        return typename(typ)
+
     def check(self) -> None:
         if self.is_internal() and self.tag is None:
             raise SerdeError('"tag" must be specified in InternalTagging')
@@ -825,39 +838,58 @@ def __call__(self, cls: T) -> _WithTagging[T]:
 
 
 @overload
-def InternalTagging(tag: str) -> Tagging: ...
+def InternalTagging(tag: str, *, tags: Optional[dict[type[Any], str]] = None) -> Tagging: ...
 
 
 @overload
-def InternalTagging(tag: str, cls: T) -> _WithTagging[T]: ...
+def InternalTagging(tag: str, cls: T, *, tags: Optional[dict[type[Any], str]] = None) -> _WithTagging[T]: ...
 
 
-def InternalTagging(tag: str, cls: Optional[T] = None) -> Union[Tagging, _WithTagging[T]]:
-    tagging = Tagging(tag, kind=Tagging.Kind.Internal)
+def InternalTagging(tag: str, cls: Optional[T] = None, *, tags: Optional[dict[type[Any], str]] = None) -> Union[Tagging, _WithTagging[T]]:
+    tagging = Tagging(tag, kind=Tagging.Kind.Internal, tags=tags)
     if cls:
         return tagging(cls)
     else:
         return tagging
 
 
 @overload
-def AdjacentTagging(tag: str, content: str) -> Tagging: ...
+def AdjacentTagging(tag: str, content: str, *, tags: Optional[dict[type[Any], str]] = None) -> Tagging: ...
 
 
 @overload
-def AdjacentTagging(tag: str, content: str, cls: T) -> _WithTagging[T]: ...
+def AdjacentTagging(tag: str, content: str, cls: T, *, tags: Optional[dict[type[Any], str]] = None) -> _WithTagging[T]: ...
 
 
 def AdjacentTagging(
-    tag: str, content: str, cls: Optional[T] = None
+    tag: str, content: str, cls: Optional[T] = None, *, tags: Optional[dict[type[Any], str]] = None
 ) -> Union[Tagging, _WithTagging[T]]:
-    tagging = Tagging(tag, content, kind=Tagging.Kind.Adjacent)
+    tagging = Tagging(tag, content, kind=Tagging.Kind.Adjacent, tags=tags)
+    if cls:
+        return tagging(cls)
+    else:
+        return tagging
+
+
+@overload
+def ExternalTagging_(*, tags: dict[type[Any], str]) -> Tagging: ...
+
+
+@overload
+def ExternalTagging_(cls: T, *, tags: dict[type[Any], str]) -> _WithTagging[T]: ...
+
+
+def ExternalTagging_(cls: Optional[T] = None, *, tags: dict[type[Any], str]) -> Union[Tagging, _WithTagging[T]]:
+    tagging = Tagging(kind=Tagging.Kind.External, tags=tags)
     if cls:
         return tagging(cls)
     else:
         return tagging
 
 
+# TODO: This is an instance rather than a function for backwards-compatibility
+# reasons. In the next major version increase this should be replaced with a
+# function.
 ExternalTagging = Tagging()
 
 Untagged = Tagging(kind=Tagging.Kind.Untagged)

serde/de.py

@@ -1129,18 +1129,18 @@ def {{func}}(cls=cls, maybe_generic=None, maybe_generic_type_vars=None, data=Non
   try:
     # create fake dict so we can reuse the normal render function
     {% if tagging.is_external() and is_taggable(t)  %}
-    ensure("{{typename(t)}}" in data , "'{{typename(t)}}' key is not present")
-    fake_dict = {"fake_key": data["{{typename(t)}}"]}
+    ensure("{{tagging.tag_for(t)}}" in data , "'{{tagging.tag_for(t)}}' key is not present")
+    fake_dict = {"fake_key": data["{{tagging.tag_for(t)}}"]}
 
     {% elif tagging.is_internal() and is_taggable(t) %}
     ensure("{{tagging.tag}}" in data , "'{{tagging.tag}}' key is not present")
-    ensure("{{typename(t)}}" == data["{{tagging.tag}}"], "tag '{{typename(t)}}' isn't found")
+    ensure("{{tagging.tag_for(t)}}" == data["{{tagging.tag}}"], "tag '{{tagging.tag_for(t)}}' isn't found")
     fake_dict = {"fake_key": data}
 
     {% elif tagging.is_adjacent() and is_taggable(t) %}
     ensure("{{tagging.tag}}" in data , "'{{tagging.tag}}' key is not present")
     ensure("{{tagging.content}}" in data , "'{{tagging.content}}' key is not present")
-    ensure("{{typename(t)}}" == data["{{tagging.tag}}"], "tag '{{typename(t)}}' isn't found")
+    ensure("{{tagging.tag_for(t)}}" == data["{{tagging.tag}}"], "tag '{{tagging.tag_for(t)}}' isn't found")
     fake_dict = {"fake_key": data["{{tagging.content}}"]}
 
     {% else %}

serde/se.py

@@ -619,16 +619,16 @@ def {{func}}(obj, reuse_instances, convert_sets, skip_none=False):
   {% for t in union_args %}
   if is_instance(obj, union_args[{{loop.index0}}]):
     {% if tagging.is_external() and is_taggable(t) %}
-    return {"{{typename(t)}}": {{rvalue(arg(t))}}}
+    return {"{{tagging.tag_for(t)}}": {{rvalue(arg(t))}}}
 
     {% elif tagging.is_internal() and is_taggable(t) %}
     res = {{rvalue(arg(t))}}
-    res["{{tagging.tag}}"] = "{{typename(t)}}"
+    res["{{tagging.tag}}"] = "{{tagging.tag_for(t)}}"
     return res
 
     {% elif tagging.is_adjacent() and is_taggable(t) %}
     res = {"{{tagging.content}}": {{rvalue(arg(t))}}}
-    res["{{tagging.tag}}"] = "{{typename(t)}}"
+    res["{{tagging.tag}}"] = "{{tagging.tag_for(t)}}"
     return res
 
     {% else %}

tests/test_union.py

@@ -27,6 +27,7 @@
     to_tuple,
     InternalTagging,
     AdjacentTagging,
+    ExternalTagging_,
     Untagged,
 )
 from serde.json import from_json, to_json
@@ -842,3 +843,41 @@ class Foo:
 
     f = Foo([10])
     assert f == from_json(Foo, to_json(f))
+
+
+T = TypeVar("T")
+
+
+def _test_union_with_custom_tags_arguments():
+    @dataclass
+    class Foo:
+        v: int
+
+    @dataclass
+    class Bar:
+        w: str
+
+    return [
+        (
+            ExternalTagging_(tags={Foo: "f", Bar: "b"})(Foo | Bar),
+            Foo(1),
+            '{"f":{"v":1}}',
+        ),
+        (
+            InternalTagging("type", tags={Foo: "f", Bar: "b"})(Foo | Bar),
+            Foo(1),
+            '{"v":1,"type":"f"}',
+        ),
+        (
+            AdjacentTagging("type", "content", tags={Foo: "f", Bar: "b"})(Foo | Bar),
+            Foo(1),
+            '{"content":{"v":1},"type":"f"}',
+        ),
+    ]
+
+
+@pytest.mark.parametrize("cls,deserialized,serialized", _test_union_with_custom_tags_arguments())
+def test_union_with_custom_tags(cls: type[T], deserialized: T, serialized: str) -> None:
+    assert to_json(deserialized, cls=cls) == serialized
+    assert from_json(cls, serialized) == deserialized
+