Source code for Muscat.IO.PickleTools

# -*- coding: utf-8 -*-
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
#
from Muscat.IO.IOFactory import RegisterWriterClass, RegisterReaderClass
from typing import Any
import pickle as __pickle

from Muscat.IO.ReaderBase import ReaderBase


class IOHelper:
    """helper Class that represent the data from a file"""

    def __init__(self, data) -> None:
        self.unnamed = data[0]
        self.named = data[1]

    def __str__(self) -> str:
        res = " named : " + str(self.named) + "\n"
        res += " unnamed : " + str(self.unnamed) + "\n"
        return res


def SaveData(filename: str, *argv, **kwargs) -> int:
    """Save the variables into the disk and return 0 if all ok

    Save variables into the disk, you can use unnamed or named variables (keyword)

    Parameters
    ----------
    filename : str
        File name to

    Returns
    -------
    int
        0 if everything ok
    """
    with open(filename, "wb") as pickle_file:
        pickler_data = __pickle.Pickler(pickle_file)
        pickler_data.dump([argv, kwargs])
        return 0
    return 1  # pragma: no cover


def LoadData(filename: str) -> Any:
    """Load data from disk using pickle format

    Load data saved with the 'saveData' from file


    Parameters
    ----------
    filename : str
        File name to use

    Returns
    -------
    IOHelper | None
        return an instance of IOHelper if ok
        return None if not ok
    """
    with open(filename, "rb") as pickle_file:
        unpickler_data = __pickle.Unpickler(pickle_file, encoding="latin1")
        data = unpickler_data.load()
        return IOHelper(data)
    return None  # pragma: no cover


class PickleReader(ReaderBase):
    """Class handling the reading of data using pickle"""

    def __init__(self, fileName: str = ""):
        super().__init__(fileName=fileName)
        self.canHandleTemporal = False

    def Read(self):
        """Reads data using pickle

        Returns
        -------
        any
            read data
        """
        internalReader = LoadData(self.fileName)
        return internalReader.named["mesh"]


class PickleWriter(object):
    """Class handling the writing of data using pickle"""

    def __init__(self):
        super().__init__()
        self.filename = ""
        self.canHandleBinaryChange = False

    def SetBinary(self, val=True):
        pass

    def SetFileName(self, filename):
        """Sets filename

        Parameters
        ----------
        filename : str
            name of the file to write
        """
        self.filename = filename

    def Open(self, fileName=None):
        if not fileName is None:
            self.SetFileName(fileName)

    def Close(self):
        pass

    def Write(self, mesh, PointFields=None, CellFields=None, GridFields=None, PointFieldsNames=None, CellFieldsNames=None, GridFieldsNames=None):
        """Writes data using pickle"""
        if PointFieldsNames is not None:
            nodeFields = {k: v for k, v in zip(PointFieldsNames, PointFields)}
        else:
            nodeFields = {}

        if CellFieldsNames is not None:
            elemFields = {k: v for k, v in zip(CellFieldsNames, CellFields)}
        else:
            elemFields = {}
        import copy

        cmesh = copy.copy(mesh)
        cmesh.nodeFields = nodeFields
        cmesh.elemFields = elemFields
        SaveData(self.filename, mesh=cmesh)


RegisterReaderClass(".pickle", PickleReader)
RegisterWriterClass(".pickle", PickleWriter)


def CheckIntegrity(GUI: bool = False):
    """AutoTest routine"""

    import numpy as np

    from Muscat.Helpers.IO.TemporaryDirectory import TemporaryDirectory

    # create a temp file
    tempdir = TemporaryDirectory.GetTempPath()

    SaveData(tempdir + "testFile.data", "two", 3, (3, 5), toto=10)
    # load data
    b = LoadData(tempdir + "testFile.data")
    # test correct data
    np.testing.assert_equal(b.unnamed[0], "two")
    np.testing.assert_equal(b.unnamed[1], 3)
    np.testing.assert_equal(b.unnamed[2], (3, 5))
    np.testing.assert_equal(b.named["toto"], 10)
    print(b)

    from Muscat.Containers.MeshCreationTools import CreateUniformMeshOfBars

    barmesh = CreateUniformMeshOfBars(0, 8, 10)

    PointFields = [np.arange(barmesh.GetNumberOfNodes())]
    PointFieldsNames = ["PointData"]

    CellFields = [np.arange(barmesh.GetNumberOfElements())]
    CellFieldsNames = ["PointData"]

    print(barmesh)
    pw = PickleWriter()
    pw.SetBinary()  # this has no effect
    pw.Open(tempdir + "testFile.pickle")
    pw.Write(barmesh)
    pw.Write(barmesh, PointFieldsNames=PointFieldsNames, PointFields=PointFields, CellFields=CellFields, CellFieldsNames=CellFieldsNames)
    pw.Close()  # this has no effect
    pr = PickleReader()
    pr.SetFileName(tempdir + "testFile.pickle")
    barmeshII = pr.Read()

    from Muscat.Containers.MeshTools import IsClose

    IsClose(barmesh, barmeshII)
    print(barmeshII)

    return "Ok"


if __name__ == "__main__":
    # import time
    # stime = time.time()
    print(CheckIntegrity())  # pragma: no cover

    # print(time.time()-stime)