From 4d09e0d2f1ce3af93479ae6b50a1609157818eb7 Mon Sep 17 00:00:00 2001
From: Langspielplatte <langspielplatte@black-mesa.xyz>
Date: Tue, 24 Jan 2023 13:51:14 +0100
Subject: [PATCH] Adds export to pickle format.

---
 plot_single_transmission_timeline.py | 5 +++--
 plot_transmission_timeline.py        | 6 ++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/plot_single_transmission_timeline.py b/plot_single_transmission_timeline.py
index 97f89b4..426bc63 100755
--- a/plot_single_transmission_timeline.py
+++ b/plot_single_transmission_timeline.py
@@ -2,6 +2,7 @@
 import math
 import multiprocessing
 import os
+import pickle
 from argparse import ArgumentParser
 
 import matplotlib
@@ -162,8 +163,8 @@ if __name__ == "__main__":
 
         if args.save:
             plt.savefig("{}{}_plot.pdf".format(args.save, csv.replace(".csv", "")))
-        else:
-            plt.show()
+        if args.export:
+            pickle.dump(fig, open("{}{}_plot.pkl".format(args.save, csv.replace(".csv", "")), "wb"))
         counter += 1
 
         plt.clf()
\ No newline at end of file
diff --git a/plot_transmission_timeline.py b/plot_transmission_timeline.py
index bd822d7..a9217b8 100755
--- a/plot_transmission_timeline.py
+++ b/plot_transmission_timeline.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import multiprocessing
 import os
+import pickle
 from argparse import ArgumentParser
 from math import ceil
 from time import sleep
@@ -68,6 +69,7 @@ if __name__ == "__main__":
     parser.add_argument("-s", "--serial_file", required=True, help="Serial csv file.")
     parser.add_argument("-p", "--pcap_csv_folder", required=True, help="PCAP csv folder.")
     parser.add_argument("--save", default=None, help="Location to save pdf file.")
+    parser.add_argument("--export", default=None, help="Export figure as an pickle file.")
     parser.add_argument(
         "-c",
         "--cores",
@@ -241,8 +243,8 @@ if __name__ == "__main__":
 
     if args.save:
         plt.savefig("{}timeline_plot.pdf".format(args.save))
-    else:
-        plt.show()
+    if args.export:
+        pickle.dump(fig, open("{}timeline_plot.pkl".format(args.export), "wb"))
 
     #goodput cdf
     plt.clf()