Track

#!/usr/bin/env python3

"""Track is a utility to track stock positions and interests.

Subcommands:
 - buy         Buy a stock
 - sell        Sell a stock
 - positions   Display current positions
 - watch       Add a stock to the watch list
 - unwatch     Remove a stock from the watch list
 - watchlist   Display the watch list
 - history     Display the history
 - pick        Add a stock to the picklist
 - unpick      Remove a stock from the watch list
 - picklist    Display the pick list
 - listall     Show all lists
"""

import argparse
import datetime
import string

import colors
import fetcher
import printing
import datastorage


class CommandDispatcher(object):
  """Performs the high-level function of each subcommand."""

  def __init__(self, datastore):
    self.datastore = datastore

  def _print_summary(self, ticker, summary):
    print('{}:'.format(ticker))
    print('  {holding} @{average_cost:.2f} (avg); '.format(**summary), end='')
    print('bought: {bought}, sold: {sold}'.format(**summary))

  def _make_position_scale(self, stoploss, price, takeprofit, purch_price):
    p  = '{:.2f}'.format(price or 0.0)
    sl = tp = 'oo'  # "infinity" lol
    p_color = colors.CYAN
    if stoploss:
      sl = '{:.2f}'.format(stoploss)
    if takeprofit:
      tp = '{:.2f}'.format(takeprofit)
    if price <= stoploss:
      sl = ''
    num_dashes = 50 - len(sl) - len(p) - len(tp)
    if not stoploss or not takeprofit:
      left = int(num_dashes/2)
      right = num_dashes - left
    else:
      total = takeprofit - stoploss
      percent_l = (price - stoploss) / total
      percent_r = (takeprofit - price) / total
      left = round(num_dashes * percent_l)
      right = round(num_dashes * percent_r)
      # Paint light-red if we've already lost money
      if price < purch_price:
        p_color = colors.LT_RED
      # Alert if we're about to get stopped out!
      if percent_l <= .2:
        p = '!{}'.format(p)
        if price <= stoploss:
          p_color = colors.BG_RED
        else:
          p_color = colors.RED
    left_bar = colors.PaintLightGray('_'*left)
    right_bar = colors.PaintLightGray('_'*right)
    bar = '{}{}{}'.format(left_bar, colors.Paint(p, p_color), right_bar)
    if stoploss:
      sl = colors.PaintRed(sl)
    if takeprofit:
      tp = colors.PaintGreen(tp)
    return '{}{}{}'.format(sl, bar, tp)

  def _format_gain_loss(self, shares, cost, current_price):
    gain_loss_per_share = current_price - cost
    gain_loss_pct = (gain_loss_per_share / current_price) * 100.0
    gain_loss_total = shares * gain_loss_per_share
    if gain_loss_pct < 0:
      color = colors.RED
    else:
      color = colors.GREEN
    txt_total = '{:,.2f}'.format(gain_loss_total)
    txt_pct = '{:.2f}%'.format(gain_loss_pct)
    return '{} ({})'.format(colors.Paint(txt_total, color), colors.Paint(txt_pct, color))

  def _print_current_positions(self):
    summaries = self.datastore.get_all_position_summaries()
    values = [(i[1]['last_update'], i[0], i[1]) for i in summaries.items()]
    values.sort()
    headers = ['TICKER', 'SHARES', 'AVG COST', 'TOTAL INVESTED', 'GAIN/LOSS', 'STOPLOSS-TAKEPROFIT SCALE']
    rows = []
    for _, ticker, summary in values:
      shares = summary['holding']
      if shares == 0:
        continue
      cost = summary['average_cost']
      stoploss = summary.get('stoploss')
      takeprofit = summary.get('takeprofit')
      current_price = fetcher.get_OHLCV(ticker)['close']
      pos_scale = self._make_position_scale(stoploss, current_price, takeprofit, cost)
      rows.append((colors.PaintBlue(ticker),
                   colors.PaintBlue(shares),
                   colors.PaintBlue('{:,.2f}'.format(cost)),
                   colors.PaintBlue('{:,.2f}'.format(shares * cost)),
                   self._format_gain_loss(shares, cost, current_price),
                   pos_scale))
    printing.TabularPrinter(headers).print(rows, indent=4, detect_pipe=False)
    return rows

  def _print_ticker_summary(self, ticker):
    summary = self.datastore.get_position_summary(ticker)
    self._print_summary(ticker, summary)

  def _print_anylist(self, name, highlights, highlight_color):
    highlight_all = '*' in highlights
    getter = getattr(self.datastore, 'get_{}list'.format(name))
    headers = ['TICKER', 'NOTE', 'AGE', 'PRICE THEN', 'PRICE NOW']
    items = [(i[1]['timestamp'], i) for i in getter().items()]
    items.sort(reverse=True)
    rows = []
    for timestamp, (ticker, data) in items:
      age = timestamp_age_string(timestamp)
      rec_price = data['prices']['close']
      price = fetcher.get_OHLCV(ticker)['close']
      pct_diff = (price - rec_price)/price*100
      price_then = '{:.2f}'.format(rec_price)
      if pct_diff < 0:
        change_disp = colors.PaintRed('{:.1f}%'.format(pct_diff))
      else:
        change_disp = colors.PaintGreen('{:.1f}%'.format(pct_diff))
      price_disp = '{:.2f}'.format(price)
      price_now = '{} ({})'.format(price_disp, change_disp)
      rows.append((colors.PaintCyan(ticker), data['note'] or '', age, price_then, price_now))
    printing.TabularPrinter(headers).print(rows, indent=4)
    return rows

  def _print_watchlist(self, highlights=[], highlight_color=colors.GREEN):
    return self._print_anylist('watch', highlights, highlight_color)

  def _print_picklist(self, highlights=[], highlight_color=colors.GREEN):
    return self._print_anylist('pick', highlights, highlight_color)

  def _print_history(self):
    for ticker,recs in self.datastore.get_history().items():
      for rec in recs:
        tt = rec[0]
        print('{}\t{}'.format(ticker, rec))

  def _print_all_lists(self, picklist_color=colors.CYAN,
                             positions_color=colors.BLUE,
                             watchlist_color=colors.YELLOW):
    print(colors.Paint('[POSITIONS]', positions_color))
    if not self._print_current_positions():
      print('    No positions')
    print(colors.Paint('\n[PICKLIST]', picklist_color))
    if not self._print_picklist(highlights=['*'], highlight_color=picklist_color):
      print('    Nothing in picklist')
    print(colors.Paint('\n[WATCHLIST]', watchlist_color))
    if not self._print_watchlist(highlights=['*'], highlight_color=watchlist_color):
      print('    Nothing in watchlist')

  def buy(self, ticker, shares, price, **kwargs):
    # selectively pluck out attributes/flags from kwargs
    attrs = {}
    for attr in ('stoploss', 'takeprofit'):
      if attr in kwargs:
        attrs[attr] = kwargs[attr]
    self.datastore.add_buy(ticker, shares, price, **attrs)
    self._print_ticker_summary(ticker)
  b = buy  # to respond to alias

  def sell(self, ticker, shares, price, **ignored):
    self.datastore.add_sell(ticker, shares, price)
    self._print_ticker_summary(ticker)
  s = sell  # to respond to alias

  def positions(self, **ignored):
    self._print_current_positions()

  def stoploss(self, ticker, price, **ignored):
    self.datastore.update_position(ticker, stoploss=price)
    self._print_current_positions()

  def takeprofit(self, ticker, price, **ignored):
    self.datastore.update_position(ticker, takeprofit=price)
    self._print_current_positions()

  def watch(self, tickers, note='', **ignored):
    for ticker in tickers:
      self.datastore.add_to_watchlist(ticker, note)
    self._print_watchlist(tickers)
  w = watch  # to respond to alias

  def unwatch(self, tickers, **ignored):
    for ticker in tickers:
      self.datastore.remove_from_watchlist(ticker)
    self._print_watchlist()

  def watchlist(self, **ignored):
    self._print_watchlist()

  def pick(self, tickers, note='', **ignored):
    for ticker in tickers:
      self.datastore.add_to_picklist(ticker, note)
    self._print_picklist(tickers)
  p = pick  # to respond to alias

  def unpick(self, tickers, **ignored):
    for ticker in tickers:
      self.datastore.remove_from_picklist(ticker)
    self._print_picklist()

  def picklist(self, **ignored):
    self._print_picklist()

  def history(self, **ignored):
    self._print_history()

  def listall(self, **ignored):
    self._print_all_lists()


def is_ticker(ticker):
  return all(c in string.ascii_uppercase for c in ticker)


def valid_ticker(ticker):
  """Determines if this string is a valid ticker and throw an exception if not.

  This is used as the type= argument to some argparse args as a quick and
  dirty way to validate arguments.
  """
  if not is_ticker(ticker):
    raise ValueError('"{}" is not a ticker'.format(ticker))
  return ticker


def timestamp_age_string(timestamp):
  """Convert a time-tuple to a human-friendly format."""
  now = datetime.datetime.now()
  then = datetime.datetime(*timestamp[:7])
  delta = now - then
  days = delta.days
  hours = int((delta.total_seconds() / 60.0 / 60.0) % 24)
  minutes = int((delta.total_seconds() / 60.0) % 60)
  if days == hours == 0:
    return '{} mins ago'.format(minutes)
  if days == 0:
    return '{}h {}m ago'.format(hours, minutes)
  return '{} days, {}:{} ago'.format(days, hours, minutes)


def main(args):
  dispatcher = CommandDispatcher(datastorage.Datastore())
  getattr(dispatcher, args.command)(**vars(args))


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  subcommands = parser.add_subparsers(title='Subcommands', dest='command')

  # buy
  parser_buy = subcommands.add_parser('buy', aliases=['b'], help='Buy a stock')
  parser_buy.add_argument('ticker')
  parser_buy.add_argument('shares', type=int)
  parser_buy.add_argument('price', type=float)
  parser_buy.add_argument('-s', '--stoploss', type=float, help='Stop-loss price.')
  parser_buy.add_argument('-t', '--takeprofit', type=float, help='Take-profit price.')

  # sell
  parser_sell = subcommands.add_parser('sell', aliases=['s'], help='Sell a stock')
  parser_sell.add_argument('ticker')
  parser_sell.add_argument('shares', type=int)
  parser_sell.add_argument('price', type=float)

  # stoploss
  parser_stoploss = subcommands.add_parser('stoploss', help='Set the stoploss price for a position')
  parser_stoploss.add_argument('ticker')
  parser_stoploss.add_argument('price', type=float)

  # takeprofit
  parser_takeprofit = subcommands.add_parser('takeprofit', help='Set the takeprofit price for a position')
  parser_takeprofit.add_argument('ticker')
  parser_takeprofit.add_argument('price', type=float)

  # positions
  parser_positions = subcommands.add_parser('positions', help='Display current positions')

  # watch
  parser_watch = subcommands.add_parser('watch', aliases=['w'], help='Add stocks to the watch list')
  parser_watch.add_argument('tickers', nargs='+', type=valid_ticker)
  parser_watch.add_argument('-n', '--note', help='')

  # unwatch
  parser_unwatch = subcommands.add_parser('unwatch', help='Remove a stock from the watch list')
  parser_unwatch.add_argument('tickers', nargs='+', type=valid_ticker)

  # watchlist
  parser_watchlist = subcommands.add_parser('watchlist', help='Display the watch list')

  # history
  parser_history = subcommands.add_parser('history', help='Display the history')

  # pick
  parser_pick = subcommands.add_parser('pick', aliases=['p'], help='Add stocks to the picklist')
  parser_pick.add_argument('tickers', nargs='+', type=valid_ticker)
  parser_pick.add_argument('-n', '--note', help='')

  # unpick
  parser_unpick = subcommands.add_parser('unpick', help='Remove a stock from the watch list')
  parser_unpick.add_argument('tickers', nargs='+', type=valid_ticker)

  # picklist
  parser_picklist = subcommands.add_parser('picklist', help='Display the pick list')

  # listall
  parser_listall = subcommands.add_parser('listall', help='Show all lists')

  args = parser.parse_args()
  if args.command == None:
    args.command = 'listall'
  main(args)