Generating deals

2 files changed, 49 insertions, 3 deletions

diff --git a/pysquaredeal.py b/pysquaredeal.py
index 246e2c6..ccfc23a 100644
--- a/pysquaredeal.py
+++ b/pysquaredeal.py
@@ -1,4 +1,4 @@
-import argparse, sys
+import argparse, os, sys
 
 from squaredeal import SquareDeal, SquareDealError
 
@@ -7,6 +7,8 @@ argparser = argparse.ArgumentParser(prog='pysquaredeal.py')
 
 argparser.add_argument('sqd_file', metavar='SQD_FILE', help='path to SQD file')
 argparser.add_argument('--sqk-file', metavar='SQK_FILE', help='path to SQK file, if not provided, deduced from SQD file', required=False)
+argparser.add_argument('--encoding', required=False, default='utf-8', metavar='ENCODING', help='SQD/SQK input file encoding, defaults to UTF-8, output is always UTF-8')
+argparser.add_argument('--bigdealx-path', required=False, metavar='BIGDEALX_PATH', help='path to bigdealx executable, defaults to BIGDEALX_PATH environment variable')
 
 subparsers = argparser.add_subparsers(title='command-specific arguments', metavar='COMMAND', dest='command')
 
@@ -24,10 +26,14 @@ argparser_phase = subparsers.add_parser('add_phase', help='add event phase')
 argparser_phase.add_argument('sessions', metavar='NO_SESSIONS', help='number of sessions in phase', type=int)
 argparser_phase.add_argument('boards', metavar='NO_BOARDS', help='number of boards in each session', type=int)
 argparser_phase.add_argument('prefix', metavar='PREFIX', help='ouput file prefix ("#" will be replaced by session number)')
-argparser_phase.add_argument('--description', required=False, metavar='DESCRIPTION', help='phase description')
+argparser_phase.add_argument('description', nargs='?', metavar='DESCRIPTION', help='phase description')
 
 argparser_publish = subparsers.add_parser('publish', help='mark SQD as published')
 
+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')
+
 arguments = argparser.parse_args()
 
 if arguments.command == 'create':
@@ -65,3 +71,12 @@ elif arguments.command == 'add_phase':
         raise SquareDealError('Cannot add phase: event already published')
     sd.add_phase(arguments.sessions, arguments.boards, arguments.prefix, description=arguments.description)
     sd.tofile(arguments.sqd_file, sqkpath=arguments.sqk_file)
+elif arguments.command == 'generate':
+    if arguments.bigdealx_path is None:
+        arguments.bigdealx_path = os.environ.get('BIGDEALX_PATH', None)
+    SquareDeal.BIGDEALX_PATH = arguments.bigdealx_path
+    sd = SquareDeal()
+    sd.fromfile(arguments.sqd_file, sqkpath=arguments.sqk_file)
+    if not sd.published:
+        raise SquareDealError('Cannot generate PBN files: event info is not marked as published')
+    sd.generate(arguments.phase, arguments.session)
diff --git a/squaredeal/__init__.py b/squaredeal/__init__.py
index 742b28d..2d4a020 100644
--- a/squaredeal/__init__.py
+++ b/squaredeal/__init__.py
@@ -1,4 +1,4 @@
-import hashlib, os, random, string, shutil
+import base64, hashlib, os, random, shutil, string, subprocess
 
 
 class SquareDealError(Exception):
@@ -26,8 +26,34 @@ class SquareDealPhase(object):
     def tostring(self):
         return ':'.join([str(self.sessions), str(self.boards), self.prefix, self.info or ''])
 
+    def _output_file_name(self, session):
+        return self.prefix.replace('#', str(session))
+
+    def generate(self, session, delayed_info, reserve=False):
+        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]
+        for session in sessions_to_generate:
+            session_key = self.s_keys[session]
+            session_key_len = int(len(session_key)/2)
+            session_left = session_key[0:session_key_len]
+            session_right = session_key[session_key_len:]
+            reserve_info = 'reserve' if reserve else 'original'
+            args = [SquareDeal.BIGDEALX_PATH,
+                    '-W', session_left,
+                    '-e', session_right,
+                    '-e', delayed_info,
+                    '-e', reserve_info,
+                    '-p', self._output_file_name(session+1),
+                    '-n', str(self.boards)]
+            subprocess.run(args, cwd=os.path.realpath(SquareDeal.BIGDEALX_PATH))
+
 
 class SquareDeal(object):
+
+    BIGDEALX_PATH=None
+
     def __init__(self):
         self.name = ''
         self.delayed_info = ''
@@ -142,3 +168,8 @@ class SquareDeal(object):
         for i in range(0, phase.sessions):
             phase.s_keys.append(self._generate_session_key())
         self.phases.append(phase)
+
+    def generate(self, phase, session):
+        phases_to_generate = range(0, len(self.phases)) if phase is None else [phase-1]
+        for phase in phases_to_generate:
+            self.phases[phase].generate(session, self.delayed_info)