update predictions cumulatively, do not delete predictions on failure

UMDBPP · Jul 24, 2021 · b35b012 · b35b012
1 parent 16b6457
commit b35b012
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/packetraven/__main__.py b/packetraven/__main__.py
@@ -393,6 +393,7 @@ def main():
         )
 
         packet_tracks = {}
+        predictions = {}
         try:
             while len(connections) > 0:
                 try:
@@ -414,13 +415,15 @@ def main():
 
                     if prediction_filename is not None:
                         try:
-                            predictions = get_predictions(
-                                packet_tracks,
-                                **{
-                                    key.replace('prediction_', ''): value
-                                    for key, value in kwargs.items()
-                                    if 'prediction_' in key
-                                },
+                            predictions.update(
+                                get_predictions(
+                                    packet_tracks,
+                                    **{
+                                        key.replace('prediction_', ''): value
+                                        for key, value in kwargs.items()
+                                        if 'prediction_' in key
+                                    },
+                                )
                             )
                             write_packet_tracks(predictions.values(), prediction_filename)
                         except PredictionError as error:

diff --git a/packetraven/gui/base.py b/packetraven/gui/base.py
@@ -766,13 +766,17 @@ async def retrieve_packets(self):
 
                 if self.toggles['prediction_file']:
                     try:
-                        self.__predictions = get_predictions(
-                            self.packet_tracks,
-                            **{
-                                key.replace('prediction_', ''): value
-                                for key, value in self.__configuration['prediction'].items()
-                                if 'prediction_' in key
-                            },
+                        self.__predictions.update(
+                            get_predictions(
+                                self.packet_tracks,
+                                **{
+                                    key.replace('prediction_', ''): value
+                                    for key, value in self.__configuration[
+                                        'prediction'
+                                    ].items()
+                                    if 'prediction_' in key
+                                },
+                            )
                         )
                         if self.prediction_filename is not None:
                             write_packet_tracks(