WIP Added support for tuples, namedtuples (both from collections and typing), sets, frozensets and OrderedDict's in MSONable, MontyEncoder and MontyDecoder

monty/collections.py

@@ -223,3 +223,8 @@ def dict2namedtuple(*args, **kwargs):
     d.update(**kwargs)
     return collections.namedtuple(typename="dict2namedtuple",
                                   field_names=list(d.keys()))(**d)
+
+
+def is_namedtuple(obj):
+    """Test if an object is a class generated from collections.namedtuple."""
+    return isinstance(obj, tuple) and hasattr(obj, "_fields") and hasattr(obj, "_asdict")

monty/json.py

@@ -6,9 +6,11 @@
 import os
 import json
 import datetime
+import sys
 
 from hashlib import sha1
 from collections import OrderedDict, defaultdict
+from collections import namedtuple
 from enum import Enum
 
 from importlib import import_module
@@ -33,6 +35,9 @@
     except ImportError:
         yaml = None  # type: ignore
 
+from monty.collections import is_namedtuple
+
+
 __version__ = "3.0.0"
 
 
@@ -62,6 +67,41 @@ def _load_redirect(redirect_file):
     return dict(redirect_dict)
 
 
+def _recursive_as_dict(obj):
+    """Recursive function to prepare serialization of objects.
+
+    Takes care of tuples, namedtuples, OrderedDict, objects with an as_dict method.
+    """
+    if is_namedtuple(obj):
+        if sys.version_info < (3, 7):  # default values for namedtuples were introduced in python 3.7.
+            return {"namedtuple_as_list": [_recursive_as_dict(it) for it in obj],
+                    "fields": obj._fields,
+                    "typename": obj.__class__.__name__,
+                    "@module": "builtins",
+                    "@class": "namedtuple"}
+        return {"namedtuple_as_list": [_recursive_as_dict(it) for it in obj],
+                "fields": obj._fields,
+                "fields_defaults": obj._fields_defaults,
+                "typename": obj.__class__.__name__,
+                "@module": "builtins",
+                "@class": "namedtuple"}
+    if isinstance(obj, tuple):
+        return {"tuple_as_list": [_recursive_as_dict(it) for it in obj],
+                "@module": "builtins",
+                "@class": "tuple"}
+    if isinstance(obj, OrderedDict):
+        return {"ordereddict_as_list": [[key, _recursive_as_dict(val)] for key, val in obj.items()],
+                "@module": "builtins",
+                "@class": "OrderedDict"}
+    if isinstance(obj, list):
+        return [_recursive_as_dict(it) for it in obj]
+    if isinstance(obj, dict):
+        return {kk: _recursive_as_dict(vv) for kk, vv in obj.items()}
+    if hasattr(obj, "as_dict"):
+        return obj.as_dict()
+    return obj
+
+
 class MSONable:
     """
     This is a mix-in base class specifying an API for msonable objects. MSON
@@ -122,15 +162,6 @@ def as_dict(self) -> dict:
 
         args = getfullargspec(self.__class__.__init__).args
 
-        def recursive_as_dict(obj):
-            if isinstance(obj, (list, tuple)):
-                return [recursive_as_dict(it) for it in obj]
-            if isinstance(obj, dict):
-                return {kk: recursive_as_dict(vv) for kk, vv in obj.items()}
-            if hasattr(obj, "as_dict"):
-                return obj.as_dict()
-            return obj
-
         for c in args:
             if c != "self":
                 try:
@@ -147,7 +178,7 @@ def recursive_as_dict(obj):
                             "a self.kwargs variable to automatically "
                             "determine the dict format. Alternatively, "
                             "you can implement both as_dict and from_dict.")
-                d[c] = recursive_as_dict(a)
+                d[c] = _recursive_as_dict(a)
         if hasattr(self, "kwargs"):
             # type: ignore
             d.update(**getattr(self, "kwargs"))  # pylint: disable=E1101
@@ -251,6 +282,8 @@ def default(self, o) -> dict:  # pylint: disable=E0202
                         "@class": "ObjectId",
                         "oid": str(o)}
 
+        # Is this still useful as we are now calling the _recursive_as_dict
+        #  method (which takes care of as_dict's) before the encoding ?
         try:
             d = o.as_dict()
             if "@module" not in d:
@@ -268,6 +301,17 @@ def default(self, o) -> dict:  # pylint: disable=E0202
         except AttributeError:
             return json.JSONEncoder.default(self, o)
 
+    def encode(self, o):
+        """MontyEncoder's encode method.
+
+        First, prepares the object by recursively transforming tuples, namedtuples,
+        object having an as_dict method and others to encodable python objects.
+        """
+        # This cannot go in the default method because default is called as a last resort,
+        # such that tuples and namedtuples have already been transformed to lists by json's encode method.
+        o = _recursive_as_dict(o)
+        return super().encode(o)
+
 
 class MontyDecoder(json.JSONDecoder):
     """
@@ -308,6 +352,21 @@ def process_decoded(self, d):
                         dt = datetime.datetime.strptime(d["string"],
                                                         "%Y-%m-%d %H:%M:%S")
                     return dt
+                if modname == "builtins":
+                    if classname == "tuple":
+                        return tuple([self.process_decoded(item) for item in d['tuple_as_list']])
+                    if classname == "namedtuple":
+                        # default values for collections.namedtuple have been introduced in python 3.7
+                        # it is probably not essential to deserialize the defaults if the object was serialized with
+                        # python >= 3.7 and deserialized with python < 3.7.
+                        if sys.version_info < (3, 7):
+                            nt = namedtuple(d['typename'], d['fields'])
+                        else:
+                            nt = namedtuple(d['typename'], d['fields'],  # pylint: disable=E1123
+                                            defaults=d['fields_defaults'])  # pylint: disable=E1123
+                        return nt(*[self.process_decoded(item) for item in d['namedtuple_as_list']])
+                    if classname == "OrderedDict":
+                        return OrderedDict([(key, self.process_decoded(val)) for key, val in d['ordereddict_as_list']])
 
                 mod = __import__(modname, globals(), locals(), [classname], 0)
                 if hasattr(mod, classname):

tests/test_collections.py

@@ -1,8 +1,11 @@
 import unittest
 import os
 
+from collections import namedtuple
+
 from monty.collections import frozendict, Namespace, AttrDict, \
     FrozenAttrDict, tree
+from monty.collections import is_namedtuple
 
 test_dir = os.path.join(os.path.dirname(__file__), 'test_files')
 
@@ -49,5 +52,27 @@ def test_tree(self):
         self.assertEqual(x['a']['b']['c']['d'], 1)
 
 
+def test_is_namedtuple():
+    a_nt = namedtuple('a', ['x', 'y', 'z'])
+    a1 = a_nt(1, 2, 3)
+    assert a1 == (1, 2, 3)
+    assert is_namedtuple(a1) is True
+    a_t = tuple([1, 2])
+    assert a_t == (1, 2)
+    assert is_namedtuple(a_t) is False
+
+    class SubList(list):
+        def _fields(self):
+            return []
+        def _fields_defaults(self):
+            return []
+        def _asdict(self):
+            return {}
+
+    sublist = SubList([3, 2, 1])
+    assert sublist == [3, 2, 1]
+    assert is_namedtuple(sublist) is False
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_json.py

@@ -7,10 +7,13 @@
 import datetime
 from bson.objectid import ObjectId
 from enum import Enum
+from collections import namedtuple
+from collections import OrderedDict
 
 from . import __version__ as tests_version
 from monty.json import MSONable, MontyEncoder, MontyDecoder, jsanitize
 from monty.json import _load_redirect
+from monty.collections import is_namedtuple
 
 test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_files")
 
@@ -188,6 +191,44 @@ def test_enum_serialization(self):
         self.assertEqual(e_new.name, e.name)
         self.assertEqual(e_new.value, e.value)
 
+    def test_tuple_serialization(self):
+        a = GoodMSONClass(a=1, b=tuple(), c=1)
+        afromdict = GoodMSONClass.from_dict(a.as_dict())
+        assert type(a.b) is tuple
+        assert type(afromdict.b) is tuple
+        assert a.b == afromdict.b
+        a = GoodMSONClass(a=1, b=[1, tuple([2, tuple([1, 2, 3, 4])])], c=1)
+        afromdict = GoodMSONClass.from_dict(a.as_dict())
+        assert type(a.b) is list
+        assert type(afromdict.b) is list
+        assert afromdict.b[0] == 1
+        assert type(afromdict.b[1]) is tuple
+        assert afromdict.b[1][0] == 2
+        assert type(afromdict.b[1][1]) is tuple
+        assert afromdict.b[1][1] == (1, 2, 3, 4)
+
+    def test_namedtuple_serialization(self):
+        a = namedtuple('A', ['x', 'y', 'zzz'])
+        b = GoodMSONClass(a=1, b=a(1, 2, 3), c=1)
+        assert is_namedtuple(b.b) is True
+        assert b.b._fields == ('x', 'y', 'zzz')
+        bfromdict = GoodMSONClass.from_dict(b.as_dict())
+        assert is_namedtuple(bfromdict.b) is True
+        assert bfromdict.b._fields == ('x', 'y', 'zzz')
+        assert bfromdict.b == b.b
+        assert bfromdict.b.x == b.b.x
+        assert bfromdict.b.y == b.b.y
+        assert bfromdict.b.zzz == b.b.zzz
+
+    def test_OrderedDict_serialization(self):
+        od = OrderedDict([('val1', 1), ('val2', 2)])
+        od['val3'] = '3'
+        a = GoodMSONClass(a=1, b=od, c=1)
+        assert list(a.b.keys()) == ['val1', 'val2', 'val3']
+        afromdict = GoodMSONClass.from_dict(a.as_dict())
+        assert type(afromdict.b) is OrderedDict
+        assert list(afromdict.b.keys()) == ['val1', 'val2', 'val3']
+
 
 class JsonTest(unittest.TestCase):
 
@@ -303,6 +344,38 @@ def test_redirect_settings_file(self):
         data = _load_redirect(os.path.join(test_dir, "test_settings.yaml"))
         self.assertEqual(data, {'old_module': {'old_class': {'@class': 'new_class', '@module': 'new_module'}}})
 
+    def test_complex_enc_dec(self):
+        a = tuple([0, 2, 4])
+        a_jsonstr = json.dumps(a, cls=MontyEncoder)
+        a_from_jsonstr = json.loads(a_jsonstr, cls=MontyDecoder)
+        assert type(a_from_jsonstr) is tuple
+
+        nt = namedtuple('A', ['x', 'y', 'zzz'])
+        nt2 = namedtuple('ABC', ['ab', 'cd', 'ef'])
+        od = OrderedDict([('val1', 1), ('val2', GoodMSONClass(a=a, b=nt2(1, 2, 3), c=1))])
+        od['val3'] = '3'
+
+        obj = nt(x=a, y=od, zzz=[1, 2, 3])
+        obj_jsonstr = json.dumps(obj, cls=MontyEncoder)
+        obj_from_jsonstr = json.loads(obj_jsonstr, cls=MontyDecoder)
+
+        assert is_namedtuple(obj_from_jsonstr) is True
+        assert obj_from_jsonstr.__class__.__name__ == 'A'
+        assert type(obj_from_jsonstr.x) is tuple
+        assert obj_from_jsonstr.x == (0, 2, 4)
+        assert type(obj_from_jsonstr.y) is OrderedDict
+        assert list(obj_from_jsonstr.y.keys()) == ['val1', 'val2', 'val3']
+        assert obj_from_jsonstr.y['val1'] == 1
+        assert type(obj_from_jsonstr.y['val2']) is GoodMSONClass
+        assert type(obj_from_jsonstr.y['val2'].a) is tuple
+        assert obj_from_jsonstr.y['val2'].a == (0, 2, 4)
+        assert is_namedtuple(obj_from_jsonstr.y['val2'].b) is True
+        assert obj_from_jsonstr.y['val2'].b.__class__.__name__ == 'ABC'
+        assert obj_from_jsonstr.y['val2'].b.ab == 1
+        assert obj_from_jsonstr.y['val2'].b.cd == 2
+        assert obj_from_jsonstr.y['val2'].b.ef == 3
+        assert obj_from_jsonstr.y['val3'] == '3'
+
 
 if __name__ == "__main__":
     unittest.main()