test_utils_serialization.py

"""
Tests for utility methods for serialization

"""

import numpy as np
import pytest

from openff.toolkit._utilities import requires_package as requires_pkg
from openff.toolkit.utils.serialization import Serializable
from openff.toolkit.utils.utils import deserialize_numpy, serialize_numpy


class Thing(Serializable):
    def __init__(self, description, mylist):
        self.description = description
        self.mylist = mylist

    def to_dict(self):
        return {"description": self.description, "mylist": self.mylist}

    @classmethod
    def from_dict(cls, d):
        return cls(d["description"], d["mylist"])

    def __eq__(self, other):
        """Comparator for asserting object field equality."""
        equality = self.description == other.description
        equality &= self.mylist == other.mylist
        return equality


# DEBUG
def write(filename, contents):
    if type(contents) is str:
        mode = "w"
    elif type(contents) is bytes:
        mode = "wb"
    else:
        raise Exception(f"Cannot handle contents of type {type(contents)}")
    with open(filename, mode) as outfile:
        outfile.write(contents)


class TestUtilsSerialization:
    """Test serialization and deserialization of a simple class."""

    @classmethod
    def setup_class(cls):
        cls.thing = Thing("blorb", [1, 2, 3])

    def test_json(self):
        """Test JSON serialization"""
        json_thing = self.thing.to_json()
        thing_from_json = self.thing.__class__.from_json(json_thing)
        assert self.thing == thing_from_json

    @requires_pkg("yaml")
    def test_yaml(self):
        """Test YAML serialization"""
        yaml_thing = self.thing.to_yaml()
        thing_from_yaml = self.thing.__class__.from_yaml(yaml_thing)
        assert self.thing == thing_from_yaml

    @requires_pkg("bson")
    def test_bson(self):
        """Test BSON serialization"""
        bson_thing = self.thing.to_bson()
        thing_from_bson = self.thing.__class__.from_bson(bson_thing)
        assert self.thing == thing_from_bson

    @pytest.mark.wip(
        reason=(
            "the current implementation of to_toml cannot handle dict "
            "keys associated to None (e.g. the ToolkitAM1BCC tag in the "
            "TestUtilsSMIRNOFFSerialization suite)."
        )
    )
    @requires_pkg("toml")
    def test_toml(self):
        """Test TOML serialization"""
        toml_thing = self.thing.to_toml()
        thing_from_toml = self.thing.__class__.from_toml(toml_thing)
        assert self.thing == thing_from_toml

    @requires_pkg("msgpack")
    def test_messagepack(self):
        """Test MessagePack serialization"""
        messagepack_thing = self.thing.to_messagepack()
        thing_from_messagepack = self.thing.__class__.from_messagepack(messagepack_thing)
        assert self.thing == thing_from_messagepack

    @pytest.mark.wip(
        reason="from/to_xml is not implemented yet. This test fails "
        "because the list of integers is saved in the XML as "
        "a list of numeric strings."
    )
    def test_xml(self):
        """Test XML serialization"""
        xml_thing = self.thing.to_xml()
        thing_from_xml = self.thing.__class__.from_xml(xml_thing)
        assert self.thing == thing_from_xml

    def test_pickle(self):
        """Test pickle serialization"""
        pkl_thing = self.thing.to_pickle()
        thing_from_pkl = self.thing.__class__.from_pickle(pkl_thing)
        assert self.thing == thing_from_pkl


class TestNumPySerialization:
    @pytest.mark.parametrize("endian", [">", "<"])
    def test_serialize_endianness(self, endian):
        """Test that arrays are serialized as big-endian"""
        # Force big- (">") or little-endianness ("<") , overriding architecture default
        dtype = np.dtype(float).newbyteorder(endian)
        arr = np.arange(3).astype(dtype)

        # This bytestring should be a big-endian representation, generated by
        # np.arange(3).astype(np.dtype(float).newbyteorder(">")).tobytes()
        assert serialize_numpy(arr)[0] == (
            b"\x00\x00\x00\x00\x00\x00\x00\x00?\xf0\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00"
        )

    @pytest.mark.parametrize("endian", [">", "<"])
    def test_deserialize_endianness(self, endian):
        """Test that arrays are deserialized as big-endian, and that
        the deserialization breaks if the input is little-endian."""
        dt_input = np.dtype(float).newbyteorder(endian)
        arr = np.arange(3, dtype=dt_input)
        deserialized = deserialize_numpy(arr.tobytes(), arr.shape)
        assert np.allclose(arr, deserialized) == (endian == ">")


class DictionaryContainer(Serializable):
    def __init__(self, dictionary):
        import copy

        self.dictionary = copy.deepcopy(dictionary)

    def to_dict(self):
        return self.dictionary

    @staticmethod
    def from_dict(dictionary):
        return DictionaryContainer(dictionary)

    def __eq__(self, other):
        """Comparator for asserting object field equality."""
        return self.dictionary == other.dictionary


class TestUtilsSMIRNOFFSerialization(TestUtilsSerialization):
    """Test serialization and deserialization of smirnoff99Frosst."""

    @classmethod
    def setup_class(cls):
        cls.thing = Thing("blorb", [1, 2, 3])
        # Create an example object holding the SMIRNOFF xmltodict dictionary representation
        import xmltodict

        from openff.toolkit.utils import get_data_file_path

        filename = get_data_file_path("test_forcefields/test_forcefield.offxml")
        with open(filename) as f:
            xml = f.read()
            dictionary = xmltodict.parse(xml)
            cls.thing = DictionaryContainer(dictionary)