JSON (de-)serialization (#124)

chatminer/chatparsers.py

@@ -35,17 +35,41 @@ def append(self, mess: ParsedMessage):
         self._parsed_messages.append(mess)
 
     def get_df(self):
-        messages_as_dict: List[Dict[str, Any]] = []
-        for mess in self._parsed_messages:
-            messages_as_dict.append(asdict(mess))
-
+        messages_as_dict = [asdict(mess) for mess in self._parsed_messages]
         df = pd.DataFrame(messages_as_dict)
         df["weekday"] = df["timestamp"].dt.day_name()
         df["hour"] = df["timestamp"].dt.hour
         df["words"] = df["message"].apply(lambda s: len(s.split(" ")))
         df["letters"] = df["message"].apply(len)
         return df
 
+    def write_to_json(self, file: str):
+        def serialize_message(mess: ParsedMessage):
+            return {
+                "timestamp": mess.timestamp.isoformat(),
+                "author": mess.author,
+                "message": mess.message,
+            }
+
+        with open(file, "w") as json_file:
+            json.dump(
+                [serialize_message(mess) for mess in self._parsed_messages],
+                json_file,
+                indent=4,
+            )
+
+    def read_from_json(self, file: str):
+        def deserialize_message(mess: dict):
+            timestamp = dt.datetime.fromisoformat(mess["timestamp"])
+            author = mess["author"]
+            message = mess["message"]
+            return ParsedMessage(timestamp=timestamp, author=author, message=message)
+
+        with open(file, "r") as json_file:
+            self._parsed_messages = [
+                deserialize_message(mess) for mess in json.load(json_file)
+            ]
+
 
 class Parser(ABC):
     def __init__(self, filepath: str):