From cf31a3a79e1dc4dd848fb2f4f4c7e8e0327bb258 Mon Sep 17 00:00:00 2001
From: Andre Merzky <andre@merzky.net>
Date: Sat, 6 Apr 2024 00:52:25 +0200
Subject: [PATCH] simpler code

---
 src/radical/utils/serialize.py | 28 +++++++++++-----------------
 src/radical/utils/typeddict.py | 30 +++++++++++++++++++-----------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/src/radical/utils/serialize.py b/src/radical/utils/serialize.py
index 55f476ff..d87969ba 100644
--- a/src/radical/utils/serialize.py
+++ b/src/radical/utils/serialize.py
@@ -3,7 +3,11 @@
 import json
 import msgpack
 
+from .typeddict import as_dict
 
+
+# ------------------------------------------------------------------------------
+#
 class _CType:
 
     def __init__(self, ctype, encode, decode):
@@ -37,15 +41,8 @@ def register_serializable(cls, encode=None, decode=None):
 # ------------------------------------------------------------------------------
 #
 def _prep_typed_dict(d):
-    from .typeddict import TypedDict
-    tmp = dict()
-    for k,v in d.items():
-        if isinstance(v, TypedDict):
-            tmp[k] = _prep_typed_dict(v)
-        else:
-            tmp[k] = v
-    tmp['_xtype'] = type(d).__name__
-    return tmp
+    from .typeddict import as_dict
+    return as_dict(d, _annotate=True)
 
 
 # ------------------------------------------------------------------------------
@@ -56,17 +53,14 @@ class _json_encoder(json.JSONEncoder):
     '''
 
     def encode(self, o, *args, **kw):
-        from .typeddict import TypedDict
-        if isinstance(o, TypedDict):
-            tmp = _prep_typed_dict(o)
-            return super().encode(tmp, *args, **kw)
-        return super().encode(o, *args, **kw)
+        tmp = as_dict(o, _annotate=True)
+        return super().encode(tmp, *args, **kw)
 
     def default(self, obj):
       # print('encode: %s' % obj)
         for cname,methods in _ctypes.items():
             if isinstance(obj, methods.ctype):
-                return {'_xtype': cname,
+                return {'_type': cname,
                         'as_str': methods.encode(obj)}
         return super().default(obj)
 
@@ -79,8 +73,8 @@ def _json_decoder(obj):
   # print('decode: %s' % obj)
     for cname, methods in _ctypes.items():
       # print('check %s' % cname)
-        if '_xtype' in obj and obj['_xtype'] == cname:
-            del obj['_xtype']
+        if '_type' in obj and obj['_type'] == cname:
+            del obj['_type']
           # print('found %s' % cname)
             if 'as_str' in obj:
                 return methods.decode(obj['as_str'])
diff --git a/src/radical/utils/typeddict.py b/src/radical/utils/typeddict.py
index fea6d21f..2173f140 100644
--- a/src/radical/utils/typeddict.py
+++ b/src/radical/utils/typeddict.py
@@ -17,10 +17,10 @@
 #   - optional runtime type checking
 #
 
+import collections
 import copy
 import sys
 
-from .serialize import register_serializable
 from .misc      import as_list, as_tuple, is_string
 
 
@@ -139,6 +139,8 @@ def __init__(self, from_dict=None, **kwargs):
               `kwargs`).
         '''
 
+        from .serialize import register_serializable
+
         register_serializable(self.__class__)
 
         self.update(copy.deepcopy(self._defaults))
@@ -315,8 +317,8 @@ def __repr__(self):
 
     # --------------------------------------------------------------------------
     #
-    def as_dict(self):
-        return as_dict(self._data)
+    def as_dict(self, _annotate=False):
+        return as_dict(self._data, _annotate)
 
 
     # --------------------------------------------------------------------------
@@ -486,21 +488,27 @@ def _query(self, key, default=None, last_key=True):
 
 # ------------------------------------------------------------------------------
 #
-def _as_dict_value(v):
-    return v.as_dict() if isinstance(v, TypedDict) else as_dict(v)
+def _as_dict_value(v, _annotate=False):
+    if isinstance(v, TypedDict):
+        ret = copy.deepcopy(v)
+        if _annotate:
+            ret['_type'] = type(v).__name__
+        return ret
+    else:
+        return as_dict(v, _annotate)
 
 
-def as_dict(src):
+def as_dict(src, _annotate=False):
     '''
     Iterate given object, apply `as_dict()` to all typed
     values, and return the result (effectively a shallow copy).
     '''
     if isinstance(src, dict):
-        tgt = {k: _as_dict_value(v) for k, v in src.items()}
-    elif isinstance(src, list):
-        tgt = [_as_dict_value(x) for x in src]
-    elif isinstance(src, tuple):
-        tgt = tuple([_as_dict_value(x) for x in src])
+        tgt = {k: _as_dict_value(v, _annotate) for k, v in src.items()}
+        if _annotate:
+            tgt['_type'] = type(src).__name__
+    elif isinstance(src, collections.abc.Iterable):
+        tgt = [_as_dict_value(x, _annotate) for x in src]
     else:
         tgt = src
     return tgt