From ec226463e392c53fbbeed709e77f676756bac060 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Wed, 6 Jul 2022 17:35:09 -0400
Subject: [PATCH] WIP tagged union message type API

---
 tractor/msg.py | 73 +++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 72 insertions(+), 1 deletion(-)

diff --git a/tractor/msg.py b/tractor/msg.py
index 9af3ccd..2a03b8f 100644
--- a/tractor/msg.py
+++ b/tractor/msg.py
@@ -40,10 +40,19 @@ Built-in messaging patterns, types, APIs and helpers.
 #   - https://jcristharif.com/msgspec/api.html#struct
 #   - https://jcristharif.com/msgspec/extending.html
 # via ``msgpack-python``:
-# - https://github.com/msgpack/msgpack-python#packingunpacking-of-custom-data-type
+# https://github.com/msgpack/msgpack-python#packingunpacking-of-custom-data-type
 
 from __future__ import annotations
+from contextlib import contextmanager as cm
 from pkgutil import resolve_name
+from typing import Union, Any
+
+
+from msgspec import Struct
+from msgspec.msgpack import (
+    Encoder,
+    Decoder,
+)
 
 
 class NamespacePath(str):
@@ -78,3 +87,65 @@ class NamespacePath(str):
             (ref.__module__,
              getattr(ref, '__name__', ''))
         ))
+
+
+# LIFO codec stack that is appended when the user opens the
+# ``configure_native_msgs()`` cm below to configure a new codec set
+# which will be applied to all new (msgspec relevant) IPC transports
+# that are spawned **after** the configure call is made.
+_lifo_codecs: list[
+    tuple[
+        Encoder,
+        Decoder,
+    ],
+] = [(Encoder(), Decoder())]
+
+
+def get_msg_codecs() -> tuple[
+    Encoder,
+    Decoder,
+]:
+    '''
+    Return the currently configured ``msgspec`` codec set.
+
+    The defaults are defined above.
+
+    '''
+    global _lifo_codecs
+    return _lifo_codecs[-1]
+
+
+@cm
+def configure_native_msgs(
+    tagged_structs: list[Struct],
+):
+    '''
+    Push a codec set that will natively decode
+    tagged structs provied in ``tagged_structs``
+    in all IPC transports and pop the codec on exit.
+
+    '''
+    global _lifo_codecs
+
+    # See "tagged unions" docs:
+    # https://jcristharif.com/msgspec/structs.html#tagged-unions
+
+    # "The quickest way to enable tagged unions is to set tag=True when
+    # defining every struct type in the union. In this case tag_field
+    # defaults to "type", and tag defaults to the struct class name
+    # (e.g. "Get")."
+    enc = Encoder()
+
+    types_union = Union[tagged_structs[0]] | Any
+    for struct in tagged_structs[1:]:
+        types_union |= struct
+
+    dec = Decoder(types_union)
+
+    _lifo_codecs.append((enc, dec))
+    try:
+        print("YOYOYOOYOYOYOY")
+        yield enc, dec
+    finally:
+        print("NONONONONON")
+        _lifo_codecs.pop()