From 2bb9ae555a7354f6f8c64ce8dc41e88f93956373 Mon Sep 17 00:00:00 2001
From: Cole Lyman <cole@colelyman.com>
Date: Mon, 30 Sep 2024 13:56:23 -0600
Subject: [PATCH] Implement a custom PlotException when a plot fails

---
 CRISPResso2/CRISPRessoCORE.py            | 4 ++++
 CRISPResso2/CRISPRessoMultiProcessing.py | 5 ++++-
 CRISPResso2/CRISPRessoPlot.py            | 1 -
 CRISPResso2/CRISPRessoShared.py          | 6 ++++++
 4 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/CRISPResso2/CRISPRessoCORE.py b/CRISPResso2/CRISPRessoCORE.py
index 8a515c59..b8a292a9 100644
--- a/CRISPResso2/CRISPRessoCORE.py
+++ b/CRISPResso2/CRISPRessoCORE.py
@@ -5174,6 +5174,10 @@ def get_scaffold_len(row, scaffold_start_loc, scaffold_seq):
         print_stacktrace_if_debug()
         error('Filtering error, please check your input.\n\nERROR: %s' % e)
         sys.exit(13)
+    except CRISPRessoShared.PlotException as e:
+        print_stacktrace_if_debug()
+        error(e)
+        sys.exit(14)
     except Exception as e:
         print_stacktrace_if_debug()
         error('Unexpected error, please check your input.\n\nERROR: %s' % e)
diff --git a/CRISPResso2/CRISPRessoMultiProcessing.py b/CRISPResso2/CRISPRessoMultiProcessing.py
index 7bcbe4de..e0693b9f 100644
--- a/CRISPResso2/CRISPRessoMultiProcessing.py
+++ b/CRISPResso2/CRISPRessoMultiProcessing.py
@@ -16,6 +16,9 @@
 import pandas as pd
 import traceback
 
+from CRISPResso2.CRISPRessoShared import PlotException
+
+
 def get_max_processes():
     return mp.cpu_count()
 
@@ -315,6 +318,6 @@ def run_plot(plot_func, plot_args, num_processes, process_futures, process_pool,
     except Exception as e:
         if halt_on_plot_fail:
             logger.critical(f"Plot error, halting execution \n")
-            raise e
+            raise PlotException(f'There was an error generating plot {plot_func.__name__}.')
         logger.warn(f"Plot error {e}, skipping plot \n")
         logger.debug(traceback.format_exc())
diff --git a/CRISPResso2/CRISPRessoPlot.py b/CRISPResso2/CRISPRessoPlot.py
index 81f69d91..a4d16495 100644
--- a/CRISPResso2/CRISPRessoPlot.py
+++ b/CRISPResso2/CRISPRessoPlot.py
@@ -2668,7 +2668,6 @@ def prep_alleles_table(df_alleles, reference_seq, MAX_N_ROWS, MIN_FREQUENCY):
     """
     dna_to_numbers={'-':0,'A':1,'T':2,'C':3,'G':4,'N':5}
     seq_to_numbers= lambda seq: [dna_to_numbers[x] for x in seq]
-
     X=[]
     annot=[]
     y_labels=[]
diff --git a/CRISPResso2/CRISPRessoShared.py b/CRISPResso2/CRISPRessoShared.py
index 78d2e052..15d1b33f 100644
--- a/CRISPResso2/CRISPRessoShared.py
+++ b/CRISPResso2/CRISPRessoShared.py
@@ -80,9 +80,15 @@ class OutputFolderIncompleteException(Exception):
 class InstallationException(Exception):
     pass
 
+
 class InputFileFormatException(Exception):
     pass
 
+
+class PlotException(Exception):
+    pass
+
+
 #########################################
 
 class StatusFormatter(logging.Formatter):