Source code for Muscat.IO.DillTools

# -*- 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 typing import Any

from Muscat.IO.ReaderBase import ReaderBase
from Muscat.IO.IOFactory import RegisterWriterClass, RegisterReaderClass


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 dill_file:
        import dill
        dill_file = dill.Pickler(dill_file)
        dill_file.dump([argv, kwargs])
        return 0
    return 1  # pragma: no cover


def LoadData(filename: str) -> Any:
    """Load data from disk using Dill 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 dill_file:
        import dill
        unpickler_data = dill.Unpickler(dill_file, encoding="latin1")
        data = unpickler_data.load()
        return IOHelper(data)
    return None  # pragma: no cover


class DillReader(ReaderBase):
    """Class handling the reading of data using Dill"""

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

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

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


class DillWriter(object):
    """Class handling the writing of data using Dill"""

    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 Dill"""
        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(".dill", DillReader)
RegisterWriterClass(".dill", DillWriter)


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

    from Muscat.Helpers.IO.TemporaryDirectory import TemporaryDirectory

    # create a temp file
    tempdir = TemporaryDirectory.GetTempPath()
    try:
        # Save data
        SaveData(tempdir + "testFile.data", "two", 3, (3, 5), toto=10)
        # load data
        b = LoadData(tempdir + "testFile.data")
        # test correct data
        if b.unnamed[0] != "two":
            raise Exception()# pragma: no cover
        if b.unnamed[1] != 3:
            raise Exception()# pragma: no cover
        if b.unnamed[2] != (3, 5):
            raise Exception()# pragma: no cover
        if b.named["toto"] != 10:
            raise Exception()# pragma: no cover
        output = b.__str__()
        print(b)
    except:  # pragma: no cover
        raise

    from Muscat.Containers.MeshCreationTools import CreateUniformMeshOfBars

    barmesh = CreateUniformMeshOfBars(0, 8, 10)
    import numpy as np

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

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

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

    from Muscat.Containers.MeshTools import IsClose

    IsClose(barmesh, barmeshII)
    print(barmeshII)

    return "Ok"


if __name__ == "__main__":
    print(CheckIntegrity())  # pragma: no cover