pymatgen.util.serialization — pymatgen 2020.7.3 documentation

# Source code for pymatgen.util.serialization

    # coding: utf-8
    # Copyright (c) Pymatgen Development Team.
    # Distributed under the terms of the MIT License.

    """
    Most features of this module has been moved to monty. Please refer to
    monty.json and monty.serialization documentation.
    """

    import json
    import functools
    import pickle

    from pymatgen.core.periodic_table import Element

    __author__ = "Shyue Ping Ong"
    __copyright__ = "Copyright 2012, The Materials Project"
    __version__ = "0.1"
    __maintainer__ = "Shyue Ping Ong"
    __email__ = "shyuep@gmail.com"
    __date__ = "Apr 30, 2012"

    def pmg_serialize(method):
        """
        Decorator for methods that add MSON serializations keys
        to the dictionary. See documentation of MSON for more details
        """

        @functools.wraps(method)
        def wrapper(*args, **kwargs):
            self = args[0]
            d = method(*args, **kwargs)
            # Add @module and @class
            d["@module"] = self.__class__.__module__
            d["@class"] = self.__class__.__name__
            return d

        return wrapper

    def json_pretty_dump(obj, filename):
        """
        Serialize obj as a JSON formatted stream to the given filename (
        pretty printing version)
        """
        with open(filename, "wt") as fh:
            json.dump(obj, fh, indent=4, sort_keys=4)

    class PmgPickler(pickle.Pickler):
        """
        Persistence of External Objects as described in section 12.1.5.1 of
        https://docs.python.org/3/library/pickle.html
        """

        def persistent_id(self, obj):
            """Instead of pickling as a regular class instance, we emit a
            persistent ID."""
            if isinstance(obj, Element):
                # Here, our persistent ID is simply a tuple, containing a tag and
                # a key
                return obj.__class__.__name__, obj.symbol
            else:
                # If obj does not have a persistent ID, return None. This means obj
                # needs to be pickled as usual.
                return None

    class PmgUnpickler(pickle.Unpickler):
        """
        Persistence of External Objects as described in section 12.1.5.1 of
        https://docs.python.org/3/library/pickle.html
        """

        def persistent_load(self, pid):
            """
            This method is invoked whenever a persistent ID is encountered.
            Here, pid is the tuple returned by PmgPickler.
            """
            try:
                type_tag, key_id = pid
            except Exception:
                # Sometimes we get a string such as ('Element', u'C') instead
                # of a real tuple. Use ast to evalute the expression (much safer
                # than eval).
                import ast
                type_tag, key_id = ast.literal_eval(pid)

            if type_tag == "Element":
                return Element(key_id)
            else:
                # Always raises an error if you cannot return the correct object.
                # Otherwise, the unpickler will think None is the object referenced
                # by the persistent ID.
                raise pickle.UnpicklingError(
                    "unsupported persistent object with pid %s" % pid)

    def pmg_pickle_load(filobj, **kwargs):
        r"""
        Loads a pickle file and deserialize it with PmgUnpickler.

        Args:
            filobj: File-like object
            **kwargs: Any of the keyword arguments supported by PmgUnpickler

        Returns:
            Deserialized object.
        """
        return PmgUnpickler(filobj, **kwargs).load()

    def pmg_pickle_dump(obj, filobj, **kwargs):
        r"""
        Dump an object to a pickle file using PmgPickler.

        Args:
            obj : Object to dump.
            fileobj: File-like object
            **kwargs: Any of the keyword arguments supported by PmgPickler
        """
        return PmgPickler(filobj, **kwargs).dump(obj)

    class SlotPickleMixin:
        """
        This mixin makes it possible to pickle/unpickle objects with __slots__
        defined.
        """

        def __getstate__(self):
            return dict(
                (slot, getattr(self, slot))
                for slot in self.__slots__ if hasattr(self, slot)
            )

        def __setstate__(self, state):
            for slot, value in state.items():
                setattr(self, slot, value)