Merge branch 'immae/parallelize' into dev

Fixes https://git.immae.eu/mantisbt/view.php?id=51

1 files changed, 149 insertions, 12 deletions

diff --git a/main.py b/main.py
index d4bab02..856d449 100644
--- a/main.py
+++ b/main.py
@@ -1,18 +1,155 @@
+from datetime import datetime
+import argparse
+import configparser
+import psycopg2
+import os
 import sys
-import helper, market
 
-args = helper.main_parse_args(sys.argv[1:])
+import market
+import portfolio
 
-pg_config, report_path = helper.main_parse_config(args.config)
+__all__ = ["make_order", "get_user_market"]
 
-for market_config, user_id in helper.main_fetch_markets(pg_config, args.user):
+def make_order(market, value, currency, action="acquire",
+        close_if_possible=False, base_currency="BTC", follow=True,
+        compute_value="average"):
+    """
+    Make an order on market
+    "market": The market on which to place the order
+    "value": The value in *base_currency* to acquire,
+             or in *currency* to dispose.
+             use negative for margin trade.
+    "action": "acquire" or "dispose".
+                "acquire" will buy long or sell short,
+                "dispose" will sell long or buy short.
+    "currency": The currency to acquire or dispose
+    "base_currency": The base currency. The value is expressed in that
+                     currency (default: BTC)
+    "follow": Whether to follow the order once run (default: True)
+    "close_if_possible": Whether to try to close the position at the end
+                         of the trade, i.e. reach exactly 0 at the end
+                         (only meaningful in "dispose"). May have
+                         unwanted effects if the end value of the
+                         currency is not 0.
+    "compute_value": Compute value to place the order
+    """
+    market.report.log_stage("make_order_begin")
+    market.balances.fetch_balances(tag="make_order_begin")
+    if action == "acquire":
+        trade = portfolio.Trade(
+                portfolio.Amount(base_currency, 0),
+                portfolio.Amount(base_currency, value),
+                currency, market)
+    else:
+        amount = portfolio.Amount(currency, value)
+        trade = portfolio.Trade(
+                amount.in_currency(base_currency, market, compute_value=compute_value),
+                portfolio.Amount(base_currency, 0),
+                currency, market)
+    market.trades.all.append(trade)
+    order = trade.prepare_order(
+            close_if_possible=close_if_possible,
+            compute_value=compute_value)
+    market.report.log_orders([order], None, compute_value)
+    market.trades.run_orders()
+    if follow:
+        market.follow_orders()
+        market.balances.fetch_balances(tag="make_order_end")
+    else:
+        market.report.log_stage("make_order_end_not_followed")
+        return order
+    market.report.log_stage("make_order_end")
+
+def get_user_market(config_path, user_id, debug=False):
+    pg_config, report_path = parse_config(config_path)
+    market_config = list(fetch_markets(pg_config, str(user_id)))[0][0]
+    return market.Market.from_config(market_config, debug=debug)
+
+def fetch_markets(pg_config, user):
+    connection = psycopg2.connect(**pg_config)
+    cursor = connection.cursor()
+
+    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 parse_config(config_file):
+    config = configparser.ConfigParser()
+    config.read(config_file)
+
+    if "postgresql" not in config:
+        print("no configuration for postgresql in config file")
+        sys.exit(1)
+
+    if "app" in config and "report_path" in config["app"]:
+        report_path = config["app"]["report_path"]
+
+        if not os.path.exists(report_path):
+            os.makedirs(report_path)
+    else:
+        report_path = None
+
+    return [config["postgresql"], report_path]
+
+def parse_args(argv):
+    parser = argparse.ArgumentParser(
+            description="Run the trade bot")
+
+    parser.add_argument("-c", "--config",
+            default="config.ini",
+            required=False,
+            help="Config file to load (default: config.ini)")
+    parser.add_argument("--before",
+            default=False, action='store_const', const=True,
+            help="Run the steps before the cryptoportfolio update")
+    parser.add_argument("--after",
+            default=False, action='store_const', const=True,
+            help="Run the steps after the cryptoportfolio update")
+    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",
+            action='append',
+            help="Do a different action than trading (add several times to chain)")
+    parser.add_argument("--parallel", action='store_true', default=True, dest="parallel")
+    parser.add_argument("--no-parallel", action='store_false', dest="parallel")
+
+    args = parser.parse_args(argv)
+
+    if not os.path.exists(args.config):
+        print("no config file found, exiting")
+        sys.exit(1)
+
+    return args
+
+def process(market_config, user_id, report_path, args):
     try:
-        user_market = market.Market.from_config(market_config, debug=args.debug)
-        helper.main_process_market(user_market, args.action, before=args.before, after=args.after)
+        market.Market\
+                .from_config(market_config, debug=args.debug, user_id=user_id, report_path=report_path)\
+                .process(args.action, before=args.before, after=args.after)
     except Exception as e:
-        try:
-            user_market.report.log_error("main", exception=e)
-        except:
-            print("{}: {}".format(e.__class__.__name__, e))
-    finally:
-        helper.main_store_report(report_path, user_id, user_market)
+        print("{}: {}".format(e.__class__.__name__, e))
+
+def main(argv):
+    args = parse_args(argv)
+
+    pg_config, report_path = parse_config(args.config)
+
+    if args.parallel:
+        import threading
+        market.Portfolio.start_worker()
+
+        for market_config, user_id in fetch_markets(pg_config, args.user):
+            threading.Thread(target=process, args=[market_config, user_id, report_path, args]).start()
+    else:
+        for market_config, user_id in fetch_markets(pg_config, args.user):
+            process(market_config, user_id, report_path, args)
+
+if __name__ == '__main__': # pragma: no cover
+    main(sys.argv[1:])