FIX: convert pickle to JSON serialization

9and3 · 9and3 · commit 8f0b06ad5552 · 2024-11-15T13:40:29.000+01:00
diff --git a/src/gh/components/DF_export_results/code.py b/src/gh/components/DF_export_results/code.py
@@ -52,6 +52,7 @@ def RunScript(self, i_dump: bool, i_export_dir: str, i_results):
         if i_dump is None or i_export_dir is None or i_results is None:
             return None
 
-        i_results.dump_pickle(i_export_dir)
+        if i_dump:
+            i_results.dump_serialization(i_export_dir)
 
         return None
diff --git a/src/gh/components/DF_import_results/code.py b/src/gh/components/DF_import_results/code.py
@@ -10,6 +10,6 @@ def RunScript(self, i_import_path: str):
         if i_import_path is None:
             return None
 
-        o_results = DFVizResults.load_pickle(i_import_path)
+        o_results = DFVizResults.load_serialization(i_import_path)
 
         return o_results
diff --git a/src/gh/diffCheck/diffCheck/df_error_estimation.py b/src/gh/diffCheck/diffCheck/df_error_estimation.py
@@ -8,7 +8,7 @@
 from datetime import datetime
 import os
 
-import pickle
+import json
 
 import numpy as np
 
@@ -21,6 +21,13 @@
 from diffCheck.df_geometries import DFAssembly
 
 
+class NumpyEncoder(json.JSONEncoder):
+    """ Special json encoder for numpy ndarray types. """
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)
+
 class DFInvalidData(Enum):
     """
     Enum to define the type of invalid data for joint or assembly analysis
@@ -36,6 +43,8 @@ class DFVizResults:
     """
     This class compiles the resluts of the error estimation into one object
     """
+    __serial_file_extenion: str = ".diffCheck"
+
     def __init__(self, assembly):
 
         self.assembly: DFAssembly = assembly
@@ -52,28 +61,46 @@ def __init__(self, assembly):
         self.distances_sd_deviation = []
         self.distances = []
 
-        self.__serial_file_extenion: str = ".diffCheck"
-
     def __repr__(self):
         return f"DFVizResults of({self.assembly})"
 
     def __getstate__(self):
         state = self.__dict__.copy()
+        if "assembly" in state and state["assembly"] is not None:
+            state["assembly"] = self.assembly.__getstate__()
         if "source" in state and state["source"] is not None:
             state["source"] = [df_cvt_bindings.cvt_dfcloud_2_dict(pcd) for pcd in state["source"]]
         if "target" in state and state["target"] is not None:
             state["target"] = [mesh.ToJSON(SerializationOptions()) for mesh in state["target"]]
+        if "sanity_check" in state and state["sanity_check"] is not None:
+            state["sanity_check"] = [s.value if isinstance(s, DFInvalidData) else s for s in self.sanity_check]
         return state
 
     def __setstate__(self, state: typing.Dict):
+        if "assembly" in state and state["assembly"] is not None:
+            assembly = DFAssembly.__new__(DFAssembly)
+            assembly.__setstate__(state["assembly"])
+            state["assembly"] = assembly
         if "source" in state and state["source"] is not None:
-            self.source = [df_cvt_bindings.cvt_dict_2_dfcloud(state["source"][i]) for i in range(len(state["source"]))]
+            source = []
+            for pcd_dict in state["source"]:
+                pcd = diffcheck_bindings.dfb_geometry.DFPointCloud()
+                pcd = df_cvt_bindings.cvt_dict_2_dfcloud(pcd_dict)
+                source.append(pcd)
+            state["source"] = source
         if "target" in state and state["target"] is not None:
-            self.target = [rg.Mesh.FromJSON(state["target"][i]) for i in range(len(state["target"]))]
+            target = []
+            for mesh_json in state["target"]:
+                mesh = rg.Mesh()
+                mesh = mesh.FromJSON(mesh_json)
+                target.append(mesh)
+            state["target"] = target
+        if "sanity_check" in state and state["sanity_check"] is not None:
+            state["sanity_check"] = [DFInvalidData(s) for s in state["sanity_check"]]
         self.__dict__.update(state)
 
-    def dump_pickle(self, dir: str) -> None:
-        """ Dump the results into a pickle file for serialization """
+    def dump_serialization(self, dir: str) -> str:
+        """ Dump the results into a JSON file for serialization """
         if not os.path.exists(os.path.dirname(dir)):
             try:
                 os.makedirs(os.path.dirname(dir))
@@ -85,21 +112,25 @@ def dump_pickle(self, dir: str) -> None:
         serial_file_path = os.path.join(dir, f"{assembly_name}_{timestamp}{self.__serial_file_extenion}")
 
         try:
-            with open(serial_file_path, "wb") as f:
-                pickle.dump(self, f)
+            with open(serial_file_path, "w") as f:
+                json.dump(self.__getstate__(), f, cls=NumpyEncoder)
         except Exception as e:
             raise e
 
+        return serial_file_path
+
     @staticmethod
-    def load_pickle(file_path: str):
-        """ Load the results from a pickle file """
+    def load_serialization(file_path: str) -> 'DFVizResults':
+        """ Load the results from a JSON file """
         if not os.path.exists(file_path):
             raise FileNotFoundError(f"File {file_path} not found")
-        if not file_path.endswith(".diffCheck"):
+        if not file_path.endswith(DFVizResults.__serial_file_extenion):
             raise ValueError(f"File {file_path} is not a valid diffCheck file")
         try:
-            with open(file_path, "rb") as f:
-                obj = pickle.load(f)
+            with open(file_path, "r") as f:
+                state = json.load(f)
+                obj = DFVizResults.__new__(DFVizResults)
+                obj.__setstate__(state)
         except Exception as e:
             raise e
         return obj
diff --git a/src/gh/diffCheck/diffCheck/df_geometries.py b/src/gh/diffCheck/diffCheck/df_geometries.py
@@ -36,6 +36,12 @@ def __post_init__(self):
 
         self.__uuid = uuid.uuid4().int
 
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state: typing.Dict):
+        self.__dict__.update(state)
+
     def __repr__(self):
         return f"Vertex: X={self.x}, Y={self.y}, Z={self.z}"
 
@@ -94,13 +100,29 @@ def __post_init__(self):
 
     def __getstate__(self):
         state = self.__dict__.copy()
+        if "all_loops" in state and state["all_loops"] is not None:
+            state["all_loops"] = [[vertex.__getstate__() for vertex in loop] for loop in state["all_loops"]]
+        # note: rg.BrepFaces cannot be serialized, so we need to convert it to a Surface >> JSON >> brep >> brepface (and vice versa)
         if "_rh_brepface" in state and state["_rh_brepface"] is not None:
-            state["_rh_brepface"] = self._rh_brepface.ToJSON(SerializationOptions())
+            state["_rh_brepface"] = self.to_brep_face().DuplicateFace(True).ToJSON(SerializationOptions())
+        return state
 
     def __setstate__(self, state: typing.Dict):
+        if "all_loops" in state and state["all_loops"] is not None:
+            all_loops = []
+            for loop_state in state["all_loops"]:
+                loop = [DFVertex.__new__(DFVertex) for _ in loop_state]
+                for vertex, vertex_state in zip(loop, loop_state):
+                    vertex.__setstate__(vertex_state)
+                all_loops.append(loop)
+            state["all_loops"] = all_loops
+        # note: rg.BrepFaces cannot be serialized, so we need to convert it to a Surface >> JSON >> brep >> brepface (and vice versa)
         if "_rh_brepface" in state and state["_rh_brepface"] is not None:
-            state["_rh_brepface"] = rg.BrepFace.FromJSON(state["_rh_brepface"])
+            state["_rh_brepface"] = rg.Surface.FromJSON(state["_rh_brepface"]).Faces[0]
         self.__dict__.update(state)
+        if self._rh_brepface is not None:
+            self.from_brep_face(self._rh_brepface, self.joint_id)
+
 
     def __repr__(self):
         return f"Face id: {(self.id)}, IsJoint: {self.is_joint} Loops: {len(self.all_loops)}"
@@ -235,6 +257,22 @@ def __post_init__(self):
         # this is an automatic identifier
         self.__uuid = uuid.uuid4().int
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if "faces" in state and state["faces"] is not None:
+            state["faces"] = [face.__getstate__() for face in self.faces]
+        return state
+
+    def __setstate__(self, state: typing.Dict):
+        if "faces" in state and state["faces"] is not None:
+            faces = []
+            for face_state in state["faces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                faces.append(face)
+            state["faces"] = faces
+        self.__dict__.update(state)
+
     def __repr__(self):
         return f"Joint id: {self.id}, Faces: {len(self.faces)}"
 
@@ -305,11 +343,56 @@ def __post_init__(self):
 
     def __getstate__(self):
         state = self.__dict__.copy()
+        if "faces" in state and state["faces"] is not None:
+            state["faces"] = [face.__getstate__() for face in self.faces]
+        if "_joint_faces" in state and state["_joint_faces"] is not None:
+            state["_joint_faces"] = [face.__getstate__() for face in state["_joint_faces"]]
+        if "_side_faces" in state and state["_side_faces"] is not None:
+            state["_side_faces"] = [face.__getstate__() for face in state["_side_faces"]]
+        if "_vertices" in state and state["_vertices"] is not None:
+            state["_vertices"] = [vertex.__getstate__() for vertex in state["_vertices"]]
+        if "_joints" in state and state["_joints"] is not None:
+            state["_joints"] = [joint.__getstate__() for joint in state["_joints"]]
         if "_center" in state and state["_center"] is not None:
             state["_center"] = self._center.ToJSON(SerializationOptions())
         return state
 
     def __setstate__(self, state: typing.Dict):
+        if "faces" in state and state["faces"] is not None:
+            faces = []
+            for face_state in state["faces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                faces.append(face)
+            state["faces"] = faces
+        if "_joint_faces" in state and state["_joint_faces"] is not None:
+            joint_faces = []
+            for face_state in state["_joint_faces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                joint_faces.append(face)
+            state["_joint_faces"] = joint_faces
+        if "_side_faces" in state and state["_side_faces"] is not None:
+            side_faces = []
+            for face_state in state["_side_faces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                side_faces.append(face)
+            state["_side_faces"] = side_faces
+        if "_vertices" in state and state["_vertices"] is not None:
+            vertices = []
+            for vertex_state in state["_vertices"]:
+                vertex = DFVertex.__new__(DFVertex)
+                vertex.__setstate__(vertex_state)
+                vertices.append(vertex)
+            state["_vertices"] = vertices
+        if "_joints" in state and state["_joints"] is not None:
+            joints = []
+            for joint_state in state["_joints"]:
+                joint = DFJoint.__new__(DFJoint)
+                joint.__setstate__(joint_state)
+                joints.append(joint)
+            state["_joints"] = joints
         if "_center" in state and state["_center"] is not None:
             state["_center"] = rg.Point3d.FromJSON(state["_center"])
         self.__dict__.update(state)
@@ -445,13 +528,58 @@ def __post_init__(self):
 
     def __getstate__(self):
         state = self.__dict__.copy()
+        if "beams" in state and state["beams"] is not None:
+            state["beams"] = [beam.__getstate__() for beam in self.beams]
         if "_mass_center" in state and state["_mass_center"] is not None:
             state["_mass_center"] = self._mass_center.ToJSON(SerializationOptions())
+        if "_all_jointfaces" in state and state["_all_jointfaces"] is not None:
+            state["_all_jointfaces"] = [face.__getstate__() for face in state["_all_jointfaces"]]
+        if "_all_sidefaces" in state and state["_all_sidefaces"] is not None:
+            state["_all_sidefaces"] = [face.__getstate__() for face in state["_all_sidefaces"]]
+        if "_all_vertices" in state and state["_all_vertices"] is not None:
+            state["_all_vertices"] = [vertex.__getstate__() for vertex in state["_all_vertices"]]
+        if "_all_joints" in state and state["_all_joints"] is not None:
+            state["_all_joints"] = [joint.__getstate__() for joint in state["_all_joints"]]
         return state
 
     def __setstate__(self, state: typing.Dict):
+        if "beams" in state and state["beams"] is not None:
+            beams = []
+            for beam_state in state["beams"]:
+                beam = DFBeam.__new__(DFBeam)
+                beam.__setstate__(beam_state)
+                beams.append(beam)
+            state["beams"] = beams
         if "_mass_center" in state and state["_mass_center"] is not None:
             state["_mass_center"] = rg.Point3d.FromJSON(state["_mass_center"])
+        if "_all_jointfaces" in state and state["_all_jointfaces"] is not None:
+            joint_faces = []
+            for face_state in state["_all_jointfaces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                joint_faces.append(face)
+            state["_all_jointfaces"] = joint_faces
+        if "_all_sidefaces" in state and state["_all_sidefaces"] is not None:
+            side_faces = []
+            for face_state in state["_all_sidefaces"]:
+                face = DFFace.__new__(DFFace)
+                face.__setstate__(face_state)
+                side_faces.append(face)
+            state["_all_sidefaces"] = side_faces
+        if "_all_vertices" in state and state["_all_vertices"] is not None:
+            vertices = []
+            for vertex_state in state["_all_vertices"]:
+                vertex = DFVertex.__new__(DFVertex)
+                vertex.__setstate__(vertex_state)
+                vertices.append(vertex)
+            state["_all_vertices"] = vertices
+        if "_all_joints" in state and state["_all_joints"] is not None:
+            joints = []
+            for joint_state in state["_all_joints"]:
+                joint = DFJoint.__new__(DFJoint)
+                joint.__setstate__(joint_state)
+                joints.append(joint)
+            state["_all_joints"] = joints
         self.__dict__.update(state)
 
     def __repr__(self):
