Allowing for X-Y range in phase/session to generate

2 files changed, 24 insertions, 5 deletions

diff --git a/pysquaredeal.py b/pysquaredeal.py
index 1b74357..346e1a6 100644
--- a/pysquaredeal.py
+++ b/pysquaredeal.py
@@ -34,8 +34,8 @@ argparser_di = subparsers.add_parser('set_dv', help='edit event delayed informat
 argparser_di.add_argument('delayed_information', metavar='DELAYED_INFO', help='value of delayed information')
 
 argparser_generate = subparsers.add_parser('generate', help='generate PBN')
-argparser_generate.add_argument('phase', nargs='?', type=int, metavar='PHASE', help='phase number, if empty, all phases will be generated')
-argparser_generate.add_argument('session', nargs='?', type=int, metavar='SESSION', help='session number, if empty, all sessions will be generated')
+argparser_generate.add_argument('phase', nargs='?', metavar='PHASE', help='phase number or range, if empty, all phases will be generated')
+argparser_generate.add_argument('session', nargs='?', metavar='SESSION', help='session number or range, if empty, all sessions will be generated')
 argparser_generate.add_argument('--reserve', action='store_true', help='generate reserve board set')
 
 arguments = argparser.parse_args()
diff --git a/squaredeal/__init__.py b/squaredeal/__init__.py
index bd10ff5..90f5f4f 100644
--- a/squaredeal/__init__.py
+++ b/squaredeal/__init__.py
@@ -4,6 +4,26 @@ import base64, hashlib, os, random, re, shutil, string, subprocess
 def generate_session_key():
     return ''.join(random.choices(string.ascii_letters + string.digits, k=60))
 
+def parse_range_str(range_str, max_count):
+    range_start = 0
+    range_end = max_count
+    if range_str:
+        try:
+            range_start = int(range_str) - 1
+            range_end = range_start + 1
+        except ValueError:
+            range_match = re.match(r'([0-9]+)-([0-9]+)', range_str)
+            if range_match:
+                range_start = int(range_match.group(1))-1
+                range_end = int(range_match.group(2))
+            else:
+                raise ValueError('Invalid range string: %s' % (range_str))
+    if range_start < 0:
+        raise ValueError('Value out of range: 0')
+    if range_end > max_count:
+        raise ValueError('Value out of range: %d' % (range_end))
+    return range(range_start, range_end)
+
 class SquareDealError(Exception):
     pass
 
@@ -43,7 +63,7 @@ class SquareDealPhase(object):
         if not SquareDeal.BIGDEALX_PATH:
             raise SquareDealError('Path to BigDeal is not set, initialize SquareDeal.BIGDEALX_PATH value')
         delayed_info = base64.b64encode(delayed_info.encode('utf-8')).decode()
-        sessions_to_generate = range(0, self.sessions) if session is None else [session-1]
+        sessions_to_generate = parse_range_str(session, self.sessions)
         for session in sessions_to_generate:
             session_key = self.s_keys[session]
             session_key_len = int(len(session_key)/2)
@@ -176,7 +196,6 @@ class SquareDeal(object):
             sqdfile.writelines(sqd_contents)
 
     def generate(self, phase, session, reserve=False):
-        # TODO: parse the funky n-m syntax for phase and session
-        phases_to_generate = range(0, len(self.phases)) if phase is None else [phase-1]
+        phases_to_generate = parse_range_str(phase, len(self.phases))
         for phase in phases_to_generate:
             self.phases[phase].generate(session, self.delayed_value, reserve)