From bc67150fc606f43d6e69186240cf55cef4361a3e Mon Sep 17 00:00:00 2001
From: rhysnewell <rhys.newell94@gmail.com>
Date: Thu, 16 Nov 2023 02:22:37 +0000
Subject: [PATCH] fix: add max retries flag

---
 flight/__init__.py        |  2 +-
 flight/flight.py          | 67 ++++++++++++++++++++++-----------------
 flight/rosella/rosella.py |  2 +-
 3 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/flight/__init__.py b/flight/__init__.py
index f49459c..51bbb3f 100644
--- a/flight/__init__.py
+++ b/flight/__init__.py
@@ -1 +1 @@
-__version__ = "1.6.1"
+__version__ = "1.6.2"
diff --git a/flight/flight.py b/flight/flight.py
index c4c788f..184cf74 100644
--- a/flight/flight.py
+++ b/flight/flight.py
@@ -395,6 +395,15 @@ def main():
         default=False,
     )
 
+    refine_options.add_argument(
+        '--max_retries',
+        help='The maximum number of times to attempt to refine a bin',
+        dest='max_retries',
+        required=False,
+        type=int,
+        default=5
+    )
+
     refine_options.set_defaults(func=refine)
 
     filter_options = subparsers.add_parser(
@@ -429,31 +438,31 @@ def main():
     # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
     # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Parsing input ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
     # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
-    if (len(sys.argv) == 2 or len(sys.argv) == 1 or sys.argv[1] == '-h'
-            or sys.argv[1] == '--help'):
-        phelp()
+    # if (len(sys.argv) == 2 or len(sys.argv) == 1 or sys.argv[1] == '-h'
+    #         or sys.argv[1] == '--help'):
+    #     phelp()
+    # else:
+    args = main_parser.parse_args()
+    time = datetime.datetime.now().strftime('%H:%M:%S %d-%m-%Y')
+
+    if args.log:
+        if os.path.isfile(args.log):
+            raise Exception("File %s exists" % args.log)
+        logging.basicConfig(
+            filename=args.log,
+            level=debug[args.verbosity],
+            format='%(asctime)s %(levelname)s: %(message)s',
+            datefmt='%m/%d/%Y %I:%M:%S %p')
     else:
-        args = main_parser.parse_args()
-        time = datetime.datetime.now().strftime('%H:%M:%S %d-%m-%Y')
-
-        if args.log:
-            if os.path.isfile(args.log):
-                raise Exception("File %s exists" % args.log)
-            logging.basicConfig(
-                filename=args.log,
-                level=debug[args.verbosity],
-                format='%(asctime)s %(levelname)s: %(message)s',
-                datefmt='%m/%d/%Y %I:%M:%S %p')
-        else:
-            logging.basicConfig(
-                level=debug[args.verbosity],
-                format='%(asctime)s %(levelname)s: %(message)s',
-                datefmt='%m/%d/%Y %I:%M:%S %p')
+        logging.basicConfig(
+            level=debug[args.verbosity],
+            format='%(asctime)s %(levelname)s: %(message)s',
+            datefmt='%m/%d/%Y %I:%M:%S %p')
 
-        logging.info("Time - %s" % (time))
+    logging.info("Time - %s" % (time))
 
 
-        args.func(args)
+    args.func(args)
 
 
 def fit(args):
@@ -609,16 +618,16 @@ def write_contig(contig, assembly, f):
     f.write(fasta)
 
 
-def phelp():
-    print("""
-Usage:
-flight [SUBCOMMAND] ..
+# def phelp():
+#     print("""
+# Usage:
+# flight [SUBCOMMAND] ..
 
-Subcommands:
-bin - Bin sets of metagenomic contigs into MAGs
-fit - Genotype variants into metagenomic strains *For use with Lorikeet*
+# Subcommands:
+# bin - Bin sets of metagenomic contigs into MAGs
+# fit - Genotype variants into metagenomic strains *For use with Lorikeet*
 
-""")
+# """)
 
 
 def str2bool(v):
diff --git a/flight/rosella/rosella.py b/flight/rosella/rosella.py
index 75ea97f..c3683ea 100755
--- a/flight/rosella/rosella.py
+++ b/flight/rosella/rosella.py
@@ -410,7 +410,7 @@ def perform_refining(self, args):
                 logging.info("Refining bins...")
                 x_min, x_max, y_min, y_max = 20, 20, 20, 20 # default plotting margins
 
-                self.slow_refine(plots, 0, 10, x_min, x_max, y_min, y_max, False,
+                self.slow_refine(plots, 0, args.max_retries, x_min, x_max, y_min, y_max, False,
                                  float(args.min_completeness), float(args.max_contamination),
                                  0, args.contaminated_only, True, False)
                 self.big_contig_filter(plots, 0, 3, x_min, x_max, y_min, y_max)