Repository: OpenFreeEnergy/gufe (language: python)

Issue:
natively key mangle common non-json-friendly objects
for example, Path, Quantity / unit


Discussion / hints:
Key mangling is insufficient:

```python
def to_dict(self):
    return {
        'here_is_a_list': [
            "containing a string",
            pathlib.Path("and/a/path.txt"),
        ]
    }
```

Mangling the key doesn't help if the key points to a container.

Fix patch:
diff --git a/gufe/custom_codecs.py b/gufe/custom_codecs.py
new file mode 100644
index 00000000..302ebb4c
--- /dev/null
+++ b/gufe/custom_codecs.py
@@ -0,0 +1,36 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/gufe
+
+# custom_codecs.py: A place to keep various custom JSONCodec instances
+
+import pathlib
+import numpy as np
+
+from gufe.custom_json import JSONCodec
+
+
+PATH_CODEC = JSONCodec(
+    cls=pathlib.Path,
+    to_dict=lambda p: {'path': str(p)},
+    from_dict=lambda dct: pathlib.Path(dct['path'])
+)
+
+BYTES_CODEC = JSONCodec(
+    cls=bytes,
+    to_dict=lambda obj: {'latin-1': obj.decode('latin-1')},
+    from_dict=lambda dct: dct['latin-1'].encode('latin-1'),
+)
+
+NUMPY_CODEC = JSONCodec(
+    cls=np.ndarray,
+    to_dict=lambda obj: {
+        'dtype': str(obj.dtype),
+        'shape': list(obj.shape),
+        'bytes': obj.tobytes()
+    },
+    from_dict=lambda dct: np.frombuffer(
+        dct['bytes'], dtype=np.dtype(dct['dtype'])
+    ).reshape(dct['shape'])
+)
+
+
diff --git a/gufe/custom_json.py b/gufe/custom_json.py
new file mode 100644
index 00000000..5149a1c9
--- /dev/null
+++ b/gufe/custom_json.py
@@ -0,0 +1,199 @@
+# This file was originally part of OpenPathSampling/SimStore, which is
+# distributed under the MIT license.
+# Portions Copyright (c) 2014-2022 the contributors to OpenPathSampling
+# Permissions are the same as those listed in the gufe LICENSE
+
+from typing import Tuple, Callable, Iterable, Dict, Any, List, Type
+import json
+import functools
+
+
+class JSONCodec(object):
+    """Custom JSON encoding and decoding for non-default types.
+
+    Parameters
+    ----------
+    cls : class
+        Class for this codec. Assumes that all subclasses should be treated
+        the same way. Can be ``None`` if ``is_my_obj`` and ``is_my_dict``
+        are given.
+    to_dict : Callable
+        method that converts the object to a dictionary
+    from_dict : Callable
+        method that restores the object based on the dictionary made by
+        to_dict
+    is_my_obj : Optional[Callable]
+        Method to determine whether the input object should be treated by
+        this encoder. Default behavior is to use ``isinstance(cls)``, and to
+        create a dict that also includes the class name and the module of
+        the object.
+    is_my_dict : Optional[Callable]
+        Method to determine whether the input dictionary should be treated
+        by this decoder. Default behavior assumes usage of the default
+        ``is_my_obj``.
+    """
+    def __init__(
+        self,
+        cls: type,
+        to_dict: Callable[[Any], Dict],
+        from_dict: Callable[[Dict], Any],
+        is_my_obj: Callable[[Any], bool] = None,
+        is_my_dict=None
+    ):
+        if is_my_obj is None:
+            is_my_obj = self._is_my_obj
+
+        if is_my_dict is None:
+            is_my_dict = self._is_my_dict
+
+        self.cls = cls
+        self.to_dict = to_dict
+        self.from_dict = from_dict
+        self.is_my_obj = is_my_obj
+        self.is_my_dict = is_my_dict
+
+    def _is_my_dict(self, dct: dict) -> bool:
+        expected = ['__class__', '__module__', ':is_custom:']
+        is_custom = all(exp in dct for exp in expected)
+        return (is_custom and dct['__class__'] == self.cls.__name__
+                and dct['__module__'] == self.cls.__module__)
+
+    def _is_my_obj(self, obj: Any) -> bool:
+        return isinstance(obj, self.cls)
+
+    def default(self, obj: Any) -> Any:
+        if self.is_my_obj(obj):
+            dct = {}
+            if self.cls:
+                dct.update({
+                    '__class__': self.cls.__name__,
+                    '__module__': self.cls.__module__,
+                    ':is_custom:': True,
+                })
+            # we let the object override __class__ and __module__ if needed
+            dct.update(self.to_dict(obj))
+            return dct
+        return obj
+
+    def object_hook(self, dct: Dict) -> Any:
+        if self.is_my_dict(dct):
+            obj = self.from_dict(dct)
+            return obj
+        return dct
+
+
+def custom_json_factory(
+    coding_methods: Iterable[JSONCodec]
+) -> Tuple[Type[json.JSONEncoder], Type[json.JSONDecoder]]:
+    """Create JSONEncoder/JSONDecoder for special types.
+
+    Factory method. Dynamically creates classes that enable all the provided
+    ``coding_methods``. Returns classes, not instances, as classes are used
+    by the ``cls`` argument in ``json.loads`` / ``json.dumps``.
+
+    Parameters
+    ----------
+    coding_methods : Iterable[JSONCodec]
+        codecs to use
+
+    Returns
+    -------
+    Tuple[Type[JSONEncoder], Type[JSONDecoder]]
+        subclasses of JSONEncoder/JSONDecoder that use support the provided
+        codecs
+    """
+    class CustomJSONEncoder(json.JSONEncoder):
+        def default(self, obj):
+            for coding_method in coding_methods:
+                # If the coding method cannot handle this object, it returns
+                # the object (unchanged). So if the object is changed, we
+                # return that.
+                result = coding_method.default(obj)
+                if result is not obj:
+                    return result
+
+            # if none of our methods are useful, use standard approaches
+            # (including providing standard error)
+            return json.JSONEncoder.default(self, obj)
+
+    class CustomJSONDecoder(json.JSONDecoder):
+        def __init__(self, *args, **kwargs):
+            # technically, JSONDecoder doesn't come with an object_hook
+            # method, which is why we pass it to super here
+            super(CustomJSONDecoder, self).__init__(
+                object_hook=self.object_hook, *args, **kwargs
+            )
+
+        def object_hook(self, dct):
+            for coding_method in coding_methods:
+                # If the coding method cannot handle this dict, it returns
+                # the dict (unchanged). So if the dict is changed, we return
+                # that.
+                result = coding_method.object_hook(dct)
+                if result is not dct:
+                    return result
+
+            # if none of our methods are useful, just return the dict
+            return dct
+
+    return (CustomJSONEncoder, CustomJSONDecoder)
+
+
+class JSONSerializerDeserializer(object):
+    """
+    Tools to serialize and deserialize objects as JSON.
+
+    This wrapper object is necessary so that we can register new codecs
+    after the original initialization.
+
+    Attributes
+    ----------
+    encoder:
+        subclass of ``JSONEncoder``; use as ``json.dumps(obj, cls=encoder)``
+    decoder:
+        subclass of ``JSONDecoder``; use as ``json.loads(string,
+        cls=decoder)``
+
+    Parameters
+    ----------
+    codecs : list of :class:`.JSONCodec`s
+        codecs supported
+    """
+    def __init__(self, codecs: Iterable[JSONCodec]):
+        self.codecs: List[JSONCodec] = []
+        for codec in codecs:
+            self.add_codec(codec)
+
+        self.encoder, self.decoder = self._set_serialization()
+
+    def _set_serialization(self) -> Tuple[Type[json.JSONEncoder],
+                                          Type[json.JSONDecoder]]:
+        encoder, decoder = custom_json_factory(self.codecs)
+        self._serializer = functools.partial(json.dumps, cls=encoder)
+        self._deserializer = functools.partial(json.loads, cls=decoder)
+        return encoder, decoder
+
+    def add_codec(self, codec: JSONCodec):
+        """Add a new codec to the supported codecs
+
+        Parameters
+        ----------
+        codec : :class:`.JSONCodec`
+            codec to add
+        """
+        if codec in self.codecs:
+            return
+
+        if codec is not None:
+            self.codecs.append(codec)
+
+        self.encoder, self.decoder = self._set_serialization()
+
+
+    def serializer(self, obj: Any) -> str:
+        """Callable that dumps to JSON"""
+        return self._serializer(obj)
+
+    def deserializer(self, string: str) -> Any:
+        """Callable to loads JSON"""
+        return self._deserializer(string)
diff --git a/gufe/tokenization.py b/gufe/tokenization.py
index c6978d82..6b495bf6 100644
--- a/gufe/tokenization.py
+++ b/gufe/tokenization.py
@@ -12,7 +12,11 @@
 import copy
 from typing import Dict, Any, Callable, Union, List, Tuple
 import weakref
+from gufe.custom_json import JSONSerializerDeserializer
+from gufe.custom_codecs import PATH_CODEC, NUMPY_CODEC, BYTES_CODEC
 
+_default_json_codecs = [PATH_CODEC, NUMPY_CODEC, BYTES_CODEC]
+JSON_HANDLER = JSONSerializerDeserializer(_default_json_codecs)
 
 # maps qualified name strings to the class
 TOKENIZABLE_CLASS_REGISTRY: Dict[Tuple[str, str], "GufeTokenizable"] = {}


Test patch:
diff --git a/gufe/tests/test_custom_json.py b/gufe/tests/test_custom_json.py
new file mode 100644
index 00000000..0d2f9887
--- /dev/null
+++ b/gufe/tests/test_custom_json.py
@@ -0,0 +1,128 @@
+# This file was originally part of OpenPathSampling/SimStore, which is
+# distributed under the MIT license.
+# Portions Copyright (c) 2014-2022 the contributors to OpenPathSampling
+# Permissions are the same as those listed in the gufe LICENSE
+
+from gufe.custom_json import (
+    JSONSerializerDeserializer, custom_json_factory, JSONCodec
+)
+from gufe.custom_codecs import PATH_CODEC, BYTES_CODEC, NUMPY_CODEC
+import json
+import pathlib
+import pytest
+
+import numpy as np
+from numpy import testing as npt
+
+
+class TestJSONSerializerDeserializer(object):
+    def test_add_codec(self):
+        # without bytes codec, can't serialize numpy
+        serialization = JSONSerializerDeserializer([NUMPY_CODEC])
+        obj = np.array([[1.0, 0.0], [2.0, 3.2]])
+        with pytest.raises(TypeError):
+            serialization.serializer(obj)
+        # add the codec and it will work
+        serialization.add_codec(BYTES_CODEC)
+        serialized = serialization.serializer(obj)
+        assert len(serialization.codecs) == 2
+        reconstructed = serialization.deserializer(serialized)
+        npt.assert_equal(obj, reconstructed)
+
+    def test_add_existing_codec(self):
+        serialization = JSONSerializerDeserializer([BYTES_CODEC])
+        assert len(serialization.codecs) == 1
+        serialization.add_codec(BYTES_CODEC)
+        assert len(serialization.codecs) == 1
+
+
+class CustomJSONCodingTest(object):
+    """Base class for testing codecs.
+
+    In ``setup()``, user must define the following:
+
+    * ``self.codec``: The codec to run
+    * ``self.objs``: A list of objects to serialize
+    * ``self.dcts``: A list of expected serilized forms of each object in
+      ``self.objs``
+    """
+    def test_default(self):
+        for (obj, dct) in zip(self.objs, self.dcts):
+            assert self.codec.default(obj) == dct
+
+    def test_object_hook(self):
+        for (obj, dct) in zip(self.objs, self.dcts):
+            assert self.codec.object_hook(dct) == obj
+
+    def _test_round_trip(self, encoder, decoder):
+        for (obj, dct) in zip(self.objs, self.dcts):
+            json_str = json.dumps(obj, cls=encoder)
+            reconstructed = json.loads(json_str, cls=decoder)
+            assert reconstructed == obj
+            json_str_2 = json.dumps(obj, cls=encoder)
+            assert json_str == json_str_2
+
+    def test_round_trip(self):
+        encoder, decoder = custom_json_factory([self.codec])
+        self._test_round_trip(encoder, decoder)
+
+    def test_not_mine(self):
+        # test that the default behavior is obeyed
+        obj = {'test': 5}
+        json_str = '{"test": 5}'
+        encoder, decoder = custom_json_factory([self.codec])
+        assert json.dumps(obj, cls=encoder) == json_str
+        assert json.loads(json_str, cls=decoder) == obj
+
+
+
+class TestNumpyCoding(CustomJSONCodingTest):
+    def setup(self):
+        self.codec = NUMPY_CODEC
+        self.objs = [np.array([[1.0, 0.0], [2.0, 3.2]]),
+                     np.array([1, 0])]
+        shapes = [[2, 2], [2,]]
+        dtypes = [str(arr.dtype) for arr in self.objs]  # may change by system?
+        byte_reps = [arr.tobytes() for arr in self.objs]
+        self.dcts = [
+            {
+                ':is_custom:': True,
+                '__class__': 'ndarray',
+                '__module__': 'numpy',
+                 'shape': shape,
+                 'dtype': dtype,
+                 'bytes': byte_rep
+            }
+            for shape, dtype, byte_rep in zip(shapes, dtypes, byte_reps)
+        ]
+
+    def test_object_hook(self):
+        # to get custom equality testing for numpy
+        for (obj, dct) in zip(self.objs, self.dcts):
+            reconstructed = self.codec.object_hook(dct)
+            npt.assert_array_equal(reconstructed, obj)
+
+    def test_round_trip(self):
+        encoder, decoder = custom_json_factory([self.codec, BYTES_CODEC])
+        for (obj, dct) in zip(self.objs, self.dcts):
+            json_str = json.dumps(obj, cls=encoder)
+            reconstructed = json.loads(json_str, cls=decoder)
+            npt.assert_array_equal(reconstructed, obj)
+            json_str_2 = json.dumps(obj, cls=encoder)
+            assert json_str == json_str_2
+
+
+class TestPathCodec(CustomJSONCodingTest):
+    def setup(self):
+        self.codec = PATH_CODEC
+        self.objs = [
+            pathlib.Path("foo/bar"),
+        ]
+        self.dcts = [
+            {
+                ":is_custom:": True,
+                "__class__": "Path",
+                "__module__": 'pathlib',
+                'path': "foo/bar"
+            }
+        ]
diff --git a/gufe/tests/test_tokenization.py b/gufe/tests/test_tokenization.py
index 56c3f7b8..a3888d47 100644
--- a/gufe/tests/test_tokenization.py
+++ b/gufe/tests/test_tokenization.py
@@ -1,10 +1,11 @@
 import pytest
 import abc
 from unittest import mock
+import json
 
 from gufe.tokenization import (
     GufeTokenizable, GufeKey, tokenize, TOKENIZABLE_REGISTRY,
-    import_qualname, get_class, TOKENIZABLE_CLASS_REGISTRY,
+    import_qualname, get_class, TOKENIZABLE_CLASS_REGISTRY, JSON_HANDLER
 )
 
 
@@ -257,3 +258,13 @@ def test_import_qualname_error_none(modname, qualname):
 def test_get_class(cls_reg):
     with mock.patch.dict("gufe.tokenization.TOKENIZABLE_CLASS_REGISTRY", cls_reg):
         assert get_class(__name__, "Outer.Inner") is Outer.Inner
+
+def test_path_to_json():
+    import pathlib
+    p = pathlib.Path("foo/bar")
+    ser = json.dumps(p, cls=JSON_HANDLER.encoder)
+    deser = json.loads(ser, cls=JSON_HANDLER.decoder)
+    reser = json.dumps(deser, cls=JSON_HANDLER.encoder)
+    assert ser == reser
+    assert p == deser
+