From b35b01214342fbf6fb111a1fef293ea4a6a2e619 Mon Sep 17 00:00:00 2001 From: Zach Date: Sat, 24 Jul 2021 08:33:43 -0400 Subject: [PATCH] update predictions cumulatively, do not delete predictions on failure --- packetraven/__main__.py | 17 ++++++++++------- packetraven/gui/base.py | 18 +++++++++++------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/packetraven/__main__.py b/packetraven/__main__.py index 69a5daee..69a4faaf 100644 --- 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 index 162161dc..a3ef6188 100644 --- 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(