1 files changed, 23 insertions, 8 deletions

diff --git a/helper.py b/helper.py
index 4b9ce0d..6d28c3f 100644
--- a/helper.py
+++ b/helper.py
@@ -24,6 +24,11 @@ def main_parse_args(argv):
     parser.add_argument("--debug",
             default=False, action='store_const', const=True,
             help="Run in debug mode")
+    parser.add_argument("--user",
+            default=None, required=False, help="Only run for that user")
+    parser.add_argument("--action",
+            default=None, required=False,
+            help="Do a different action than trading")
 
     args = parser.parse_args(argv)
 
@@ -51,21 +56,31 @@ def main_parse_config(config_file):
 
     return [config["postgresql"], report_path]
 
-def main_fetch_markets(pg_config):
+def main_fetch_markets(pg_config, user):
     connection = psycopg2.connect(**pg_config)
     cursor = connection.cursor()
 
-    cursor.execute("SELECT config,user_id FROM market_configs")
+    if user is None:
+        cursor.execute("SELECT config,user_id FROM market_configs")
+    else:
+        cursor.execute("SELECT config,user_id FROM market_configs WHERE user_id = %s", user)
 
     for row in cursor:
         yield row
 
-def main_process_market(user_market, before=False, after=False):
-    if before:
-        process_sell_all__1_all_sell(user_market)
-    if after:
-        portfolio.Portfolio.wait_for_recent(user_market)
-        process_sell_all__2_all_buy(user_market)
+def main_process_market(user_market, action, before=False, after=False):
+    if action is None:
+        if before:
+            process_sell_all__1_all_sell(user_market)
+        if after:
+            portfolio.Portfolio.wait_for_recent(user_market)
+            process_sell_all__2_all_buy(user_market)
+    elif action == "print_balances":
+        print_balances(user_market)
+    elif action == "print_orders":
+        print_orders(user_market)
+    else:
+        raise NotImplementedError("Unknown action {}".format(action))
 
 def main_store_report(report_path, user_id, user_market):
     try: