"""
This is a module that extends Python's standard builtin json module. It offers
the following additional features:
- `JsonableClassMixin`, a mixin that helps with creating serializable classes
- An extended JSON encoder that can handle classes derived from `JsonableClassMixin`
- `dump` and `dumps` functions that use the extended encoder by default.
"""
import typing
import copy
import json as builtin_json
from json import JSONDecodeError # noqa: F401
from schrodinger import get_mmshare_version
from schrodinger.infra import util
NATIVE_JSON_DATATYPES = (str, int, float, bool, type(None), dict, list)
VERSION_KEY = '_version_'
VALUE_KEY = '_value_'
# Attribute name to store the version of an object when calling `.toJson()`
_TO_VERSION_ATTRNAME = '_TO_VERSION_'
# Attribute name to store the version of an object when calling `.fromJson()`
_FROM_VERSION_ATTRNAME = '_FROM_VERSION_'
[docs]class JSONEncodeError(ValueError):
pass
# We create simple subclasses of all the native json datatypes so we can
# attach version attributes to them.
_attributable_types = {}
for type_ in NATIVE_JSON_DATATYPES:
if type_ is bool or type_ is type(None):
# bool and NoneType can't be subclassed, so we don't have to worry
# about adding versions to them.
continue
[docs] class AttributableType(type_):
pass
_attributable_types[type_] = AttributableType
def _unwrap_attributable_type(attributable_inst):
orig_dataclass = type(attributable_inst).mro()[1]
return orig_dataclass(attributable_inst)
def _add_version_attribute(json_obj, version, attr_name, recursive=False):
"""
Create a new wrapped version of `json_obj` that has the attribute
`attr_name` set to the `version`. If `recursive` is set to True
and the `json_obj` is a list or dict, then all values of the `json_obj`
will also have a version attribute.
.. NOTE:
Booleans and `None`s can't be subclassed and it doesn't really
make sense to note versions on them anyways since they're singletons,
so we just return them unmodified.
:param json_obj: A json data structure
:type json_obj: An instance of one of the `NATIVE_JSON_DATATYPES`
:param version: The version of the structure, typically the mmshare
version.
:param attr_name: The attribute name to set on `json_obj`
:type attr_name: str
:param recursive: Whether to recursively set the version attribute.
:type recursive: bool
"""
if type(json_obj) is bool or json_obj is None:
return json_obj
AttributableType = _attributable_types[type(json_obj)]
versioned_obj = AttributableType(json_obj)
setattr(versioned_obj, attr_name, version)
if recursive and isinstance(versioned_obj, (dict, list)):
if isinstance(versioned_obj, dict):
items = versioned_obj.items()
else:
items = enumerate(versioned_obj)
for key, value in items:
versioned_obj[key] = _add_version_attribute(value,
version,
attr_name,
recursive=True)
return versioned_obj
[docs]class JsonableClassMixin:
"""
A mixin that aids in making a class Jsonable. Users must define
`toJsonImplementation(self)` and `fromJsonImplementation(cls, json_dict)`.
For example::
class Person(JsonableClassMixin, object):
def __init__(self, full_name, age):
self.full_name = full_name
self.age = age
def toJsonImplementation(self):
return {'full_name':self.full_name,
'age':self.age}
@classmethod
def fromJsonImplementation(cls, json_dict):
return cls(json_dict['full_name'], json_dict['age'])
Now `dump` and `dumps` can encode :code:`Person`::
# Encode to a file
abe = Person('Abraham Lincoln', 208)
with open('abe.json', 'w') as out_file:
json.dump(abe, out_file)
# Encode to a string
abe_json_string = json.dumps(abe)
If you want to decode the json string or file, you'll have to use
`loads` or `load` and then feed in the result to your class' public
class method `fromJson()`::
# Loading an object from a json file
with open('abe.json', 'r') as in_file:
abe = json.load(in_file, DataClass=Person)
# Loading an object from a json string
abe = json.loads(abe_json_string, DataClass=Person)
"""
# TODO: Add class constant that chooses whether to enforce strict json key
# conversion.
def __init_subclass__(cls):
if 'toJson' in cls.__dict__ or 'fromJson' in cls.__dict__:
err_msg = ('toJson() and fromJson() are not intended to be '
'overridden. To customize toJson() behavior, override '
'toJsonImplementation(); to customize fromJson() '
'behavior, override fromJsonImplementation().')
raise TypeError(err_msg)
cls._setJsonAdapters()
super().__init_subclass__()
[docs] def toJsonImplementation(self):
"""
Abstract method that must be defined by all derived classes. Converts
an instance of the derived class into a jsonifiable object.
:return: A dict made up of JSON native datatypes or Jsonable objects.
See the link below for a table of such types.
https://docs.python.org/2/library/json.html#encoders-and-decoders
"""
raise NotImplementedError()
[docs] @classmethod
def fromJsonImplementation(cls, json_dict):
"""
Abstract method that must be defined by all derived classes. Takes
in a dictionary and constructs an instance of the derived class.
:param json_dict: A dictionary loaded from a JSON string or file.
:type json_dict: dict
:return: An instance of the derived class.
:rtype: cls
"""
raise NotImplementedError()
[docs] def toJson(self, _mark_version=True):
"""
Create and returns a data structure made up of jsonable items.
:rtype: An instance of one the classes from NATIVE_JSON_DATATYPES
"""
d = self.toJsonImplementation()
value = _jsonify(d)
if _mark_version:
return _add_version_attribute(value,
version=self.get_version(),
attr_name=_TO_VERSION_ATTRNAME)
else:
return value
[docs] @classmethod
def fromJson(cls, json_obj):
"""
A factory method which constructs a new object from a given dict
loaded from a json string or file.
:param json_obj: A json-loaded dictionary to create an object from.
:type json_obj: dict
:return: An instance of this class.
:rtype: cls
"""
if hasattr(json_obj, _FROM_VERSION_ATTRNAME):
# If the attribute storing the version is present, then
# we're at the top level of the decoding. We unwrap
# the object into its normal datatype and then save
# the version as an attribute of this function.
version = getattr(json_obj, _FROM_VERSION_ATTRNAME)
json_obj = _unwrap_attributable_type(json_obj)
cls.fromJson.__func__.version = version
else:
# If the attribute isn't present, then check on this function
# to see if its been stored already.
version = cls.fromJson.__func__.version
if version is not None:
adapters = cls._getJsonAdapters()
for ad in adapters:
if ad.json_adapter_version > version:
json_obj = ad(json_obj)
else:
json_obj = copy.deepcopy(json_obj)
try:
cls.fromJson.__func__._stack_count += 1
return cls.fromJsonImplementation(json_obj)
finally:
cls.fromJson.__func__._stack_count -= 1
if cls.fromJson.__func__._stack_count == 0:
cls.fromJson.__func__.version = None
# Save the version of the json object we're processing as an
# attribute on the function. This should be fine since there shouldn't
# ever be an instance where two separate objects are being deserialized
# simultaneously.
fromJson.__func__.version = None
# Use a counter for keeping track of recursion depth so we don't reset
# `fromJson.__func__.version` before all child objects get a chance to be
# deserialized.
fromJson.__func__._stack_count = 0
@classmethod
def _getJsonAdapters(cls):
"""
:return: A list of all adapter methods for this class.
:rtype: list(function)
"""
json_adapters_attr_name = '_' + cls.__name__ + '_jsonAdapters'
adapters = getattr(cls, json_adapters_attr_name, None)
if adapters is None:
cls._setJsonAdapters()
return getattr(cls, json_adapters_attr_name)
@classmethod
def _setJsonAdapters(cls):
"""
Sets the class attribute storing a list of all json adapter methods.
"""
json_adapters_attr_name = '_' + cls.__name__ + '_jsonAdapters'
adapters = util.find_decorated_methods(cls, 'json_adapter_version')
# Ensure that the versions are comparable
if len({type(adapter.json_adapter_version) for adapter in adapters
}) > 1:
raise TypeError('Adapter versions must all be of the same type')
# Sort the adapters
adapters.sort(key=lambda a: a.json_adapter_version)
setattr(cls, json_adapters_attr_name, adapters)
[docs] def get_version(self):
"""
Method to get the version of a particular object. Defaults to
the current version of mmshare. This class can be overridden for custom
versioning behavior.
"""
return get_mmshare_version()
[docs]def adapter(version):
"""
This function is a decorator used to define an adapter function that takes
in a `json_obj` from an older version of a `JsonableClassMixin` and
returns it modified such that it can be read in by the version of the class
specified by the `version` argument (typically the current version at the
time the adapter is written).
As a example, imagine we define a simple class::
class Person:
# mmshare version: 40000
def __init__(self, full_name):
self.full_name = full_name
If, in mmshare version 41000, we split :code:`full_name` into attributes
:code:`first_name` and :code:`last_name`, we could define an adapter like so::
class Person:
# mmshare version: 41000
def __init__(self, first_name, last_name):
self.first_name = first_name
self.last_name = last_name
@json.adapter(version=41000)
def _JsonAdapter(self, json_dict):
full_name = json_dict.pop(full_name)
fname, lname = full_name.split(' ')
json_dict['first_name'] = fname
json_dict['last_name'] = lname
return json_dict
.. note::
An adapter function only needs to bring the `json_dict` up-to-date with
the class *at the time that the adapter is written*. The next time the
class is changed in a way that breaks json decoding, a new adapter
function should be added to the class that takes in the previous
adapter's output and makes it compatible with the current class. In
this way the json framework can decode any older version by
automatically passing it through the appropriate chain of adapter
functions.
:param version: The data version to which an older `json_dict` will be
adapted (typically the current version when the adapter is writte).
:type version: int, str, or list of str or ints
:raises TypeError: if `version` is not an int, str, or list
"""
if type(version) not in {int, str, list}:
raise TypeError('Adapter versions must an int, str, or list')
if type(version) is list and any(
not isinstance(item, (int, str)) for item in version):
raise TypeError(
'Adapter version lists must be made up of only ints or strings')
def wrapped_adapter(f):
f.json_adapter_version = version
return classmethod(f)
return wrapped_adapter
def _dict_to_json(d):
new_d = {}
for k, v in d.items():
if not isinstance(k, str):
raise JSONEncodeError(
'The JSON protocol only supports using strings '
'as keys of dictionaries')
new_d[k] = _jsonify(v)
return new_d
def _list_to_json(lst):
new_l = []
for idx, v in enumerate(lst):
new_l.append(_jsonify(v))
return new_l
# JsonableClassMixin is defined in this module, so we can't define this
# constant at the top.
JSONABLE_DATATYPES = NATIVE_JSON_DATATYPES + (JsonableClassMixin,)
def _jsonify(val):
"""
Convert `val` into a jsonable data structure. `val` must either subclass
`JsonableClassMixin` or be a jsonable data type. If `val` is a dict or
list, or `val.toJson()` returns a dict or list, then this will recursively
go through the dict or list values and jsonify them.
:return: A jsonable data structure.
:raises JsonEncodeError: if val is not jsonable.
"""
if (not isinstance(val, JsonableClassMixin) and
type(val) in jsonable.DATACLASS_REGISTRY):
serializer = jsonable.get_default_serializer(type(val))
val = serializer.jsonableFromObject(val)
if not isinstance(val, JSONABLE_DATATYPES):
# TODO: This ideally causes an error on class definition or
# instantiation, similar to an ABCMethod
JSONABLE_DATATYPE_NAMES = ' '.join(
cls.__name__ for cls in JSONABLE_DATATYPES)
err_msg = (
f"type <{type(val).__name__}> is not jsonable. To be jsonable, "
"an object must be an instance of one of the following types: " +
JSONABLE_DATATYPE_NAMES)
raise JSONEncodeError(err_msg)
if isinstance(val, JsonableClassMixin):
val = val.toJson(_mark_version=False)
if isinstance(val, dict):
return _dict_to_json(val)
elif isinstance(val, list):
return _list_to_json(val)
else:
return val
[docs]def get_json_version_from_file(json_fh):
"""
Get the version information from a json file serialized from
`JsonableClassMixin` objects. Returns None if no version information is
found.
:type json_fh: file handle open for reading
"""
json_obj = load(json_fh)
return getattr(json_obj, _FROM_VERSION_ATTRNAME, None)
[docs]def get_json_version_from_string(json_str):
"""
Get the version information from a json string serialized from
`JsonableClassMixin` objects. Returns None if no version information is
found.
"""
json_obj = loads(json_str)
return getattr(json_obj, _FROM_VERSION_ATTRNAME, None)
################################################################################
# Wrappers for the standard libs json module
################################################################################
# Import down here to prevent circular import
from . import jsonable # isort:skip
def _encode(obj, serializer=None):
"""
If the object implements JsonableClassMixin, Wrap the object in a
dictionary mapping json.VALUE_KEY to the object and json.VERSION_KEY to the
version of the object. Otherwise, use the given serializer to turn the
object into a serializable version of the object. If neither of these
conditions are met, simply return the object.
"""
# If the object hasn't already been converted into a jsonable
# structure and implements JsonableClassMixin, then call toJson.
if (not isinstance(obj, JsonableClassMixin) and
type(obj) in jsonable.DATACLASS_REGISTRY and serializer is None):
serializer = jsonable.get_default_serializer(type(obj))
if serializer is not None:
obj = serializer.jsonableFromObject(obj)
if isinstance(obj, JsonableClassMixin):
obj = obj.toJson()
elif isinstance(obj,
(dict, list)) and not hasattr(obj, _TO_VERSION_ATTRNAME):
obj = _jsonify(obj)
# If the object has already had `toJson()` called on it, it'll
# have the version attached to it.
if hasattr(obj, _TO_VERSION_ATTRNAME):
version = getattr(obj, _TO_VERSION_ATTRNAME)
obj = {VERSION_KEY: version, VALUE_KEY: obj}
return obj
[docs]def decode(obj, DataClass=None, serializer=None):
"""
Decode the `obj` from a json structure to the original object that encoded
it. This is done using the following methods in order of priority:
1) Detecting whether the jsoned object originally implemented
`JsonableClassMixin`. If it was, decode it into a normal json object
and retrieve the version number.
2) Using `DataClass` if it's a `typing` annotation generic (e.g.
`typing.List`, `typing.Tuple`, etc)
3) Using the supplied `serializer`
4) Inferring a serializer using `DataClass`
"""
# Detect whether the object was originally a `JsonableClassMixin` type
# by checking whether it's a dict that has a key VERSION_KEY
if isinstance(obj, dict) and VERSION_KEY in obj:
version = obj[VERSION_KEY]
obj = obj[VALUE_KEY]
obj = _add_version_attribute(obj,
version,
_FROM_VERSION_ATTRNAME,
recursive=False)
if DataClass is not None:
if DataClass is type(obj):
return obj
elif isinstance(DataClass, type) and issubclass(DataClass,
JsonableClassMixin):
return DataClass.fromJson(obj)
elif typing.get_origin(DataClass) in {set, list, dict, tuple}:
obj = _decode_generic(obj, DataClass)
else:
serializer = jsonable.get_default_serializer(DataClass)
if serializer is not None:
return serializer.objectFromJson(obj)
return obj
def _decode_generic(obj, generic_class):
generic = typing.get_origin(generic_class)
generic_args = typing.get_args(generic_class)
if generic is tuple:
if generic_args:
if len(generic_args) != len(obj) - 1:
raise TypeError(
f"{len(obj)-1} values specified for {generic_class}")
return tuple(
decode(v, DataClass=arg) for v, arg in zip(obj, generic_args))
else:
return decode(obj, DataClass=tuple)
elif generic is list:
if generic_args:
list_member_type = generic_args[0]
return [decode(v, DataClass=list_member_type) for v in obj]
else:
return decode(obj, DataClass=list)
elif generic is set:
if generic_args:
list_member_type = generic_args[0]
set_values = [
decode(v, DataClass=list_member_type) for v in obj[:-1]
]
else:
set_values = decode(obj, DataClass=set)
return set(set_values)
elif generic is dict:
if generic_args:
key_type = generic_args[0]
value_type = generic_args[1]
return {
decode(k, DataClass=key_type): decode(v, DataClass=value_type)
for k, v in obj.items()
}
else:
return obj
[docs]def dumps(obj, serializer=None, **kwargs):
"""
A wrapper that automatically encodes objects derived from
`JsonableClassMixin`.
:param serializer: A custom serializer to use for serializing the object.
This should not be used at the same time as the `DataClass` argument.
:type serializer: jsonable.AbstractJsonSerializer
If `obj` does not subclass `JsonableClassMixin`, then this function
will behave exactly like the builtin `json.dumps`.
"""
obj = _encode(obj, serializer)
return builtin_json.dumps(obj, **kwargs)
[docs]def dump(obj, fp, serializer=None, **kwargs):
"""
A wrapper that automatically encodes objects derived from
`JsonableClassMixin`.
:param serializer: A custom serializer to use for serializing the object.
This should not be used at the same time as the `DataClass` argument.
:type serializer: jsonable.AbstractJsonSerializer
If `obj` does not subclass `JsonableClassMixin`, then this function
will behave exactly like the builtin `json.dump`.
"""
obj = _encode(obj, serializer)
return builtin_json.dump(obj, fp, **kwargs)
[docs]def loads(json_str, DataClass=None, serializer=None, **kwargs):
"""
A wrapper that automatically decodes json strings serialized from
`JsonableClassMixin` objects.
:param DataClass: The class of the object that was serialized into
`json_str`. The class must either be in JSONABLE_DATATYPES or a type
that's been registered in `schrodinger.models.jsoable`.
:param serializer: A custom serializer to use for deserializing the string.
This should not be used at the same time as the `DataClass` argument.
:type serializer: jsonable.AbstractJsonSerializer
If the json string was not encoded `JsonableClassMixin` object, this
function will behave exactly like the builtin `json.loads`.
"""
obj = builtin_json.loads(json_str, **kwargs)
obj = decode(obj, DataClass, serializer)
return obj
[docs]def load(fp, DataClass=None, serializer=None, **kwargs):
"""
A wrapper that automatically decodes json files serialized from
`JsonableClassMixin` objects.
:param DataClass: The class of the object that was serialized into
`fp`. The class must either be in JSONABLE_DATATYPES or a type
that's been registered in `schrodinger.models.jsonable`.
:param serializer: A custom serializer to use for deserializing the file.
This should not be used at the same time as the `DataClass` argument.
:type serializer: jsonable.AbstractJsonSerializer
If the json file was not encoded `JsonableClassMixin` object, this
function will behave exactly like the builtin `json.load`.
"""
obj = builtin_json.load(fp, **kwargs)
obj = decode(obj, DataClass, serializer)
return obj
JSONDecoder = builtin_json.JSONDecoder
JSONEncoder = builtin_json.JSONEncoder