stop

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import signal
import sys

from configparser import ConfigParser
from itertools import cycle
from queue import Empty, Queue
from threading import Event, Lock, Thread

try:
    import yfinance

    HAS_YFINANCE = True
except ImportError:
    HAS_YFINANCE = False


__program__ = os.path.basename(sys.argv[0])
__version__ = "0.2.8-dev"  # major.minor.revision-channel
__author__ = "Axel McLaren <scm@axml.uk>"


_DEFAULT_CONFIG = os.path.expanduser("~/.config/stop.conf")
_DEFAULT_DELAY = 5.0  # seconds
_DEFAULT_ROUNDING = 2
_DEFAULT_PADDING = 4
_DEFAULT_COLOR = {
    'background': '#222222',
    'foreground': '#e7e7e7',
    'highlight': '#000000',
    'gain': '#55aa55',
    'loss': '#ff5555'
}
_DEFAULT_OPENER = "xdg-open"
_DEFAULT_THREADS = 8

_BASE_WIDTH = 24  # 7-char ticker + 6.2-digit value + 3.2-digit diff + 2 thousand separators
_CONCEAL_CHAR = '•'

_CONFIG_FORMAT = {
    'general': {
        'percent': bool,
        'delay': float,
        'thousand_separator': bool,
        'opener': str,
        'plain': bool,
        'rounding': int,
        'padding': int,
        'rotate': bool,
        'one_shot': bool,
        'snapshot': bool,
        'starting_ticker': str,
        'threads': int,
        'zoom_level': str,
    },
    'colors': {
        'background': str,
        'foreground': str,
        'highlight': str,
        'gain': str,
        'loss': str,
    },
    'tickers': (),
}

_FORMAT_PLAIN_VAL = "{} {} "
_FORMAT_POLYBAR_VAL = "%{{B{}}}%{{F{}}}%{{u{}}}%{{+u}}%{{o{}}}%{{+o}} {} %{{-o}}%{{-u}}%{{F-}}%{{B-}}%{{B{}}} {} "
_FORMAT_POLYBAR_URL_VAL = "%{{A2:{} \"{}\":}}%{{B{}}}%{{F{}}}%{{u{}}}%{{+u}}%{{o{}}}%{{+o}} {} %{{-o}}%{{-u}}%{{F-}}%{{B-}}%{{B{}}} {} "
_ZOOM_LEVELS = (
    # day
    {
        'level': 'day',
        'history': ('5d', '1d'),
        'shift': 1,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + "[{}]",
            'polybar': _FORMAT_POLYBAR_VAL + "[%{{F{}}}{}%{{F-}}]%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + "[%{{F{}}}{}%{{F-}}]%{{B-}}%{{A}}",
        },
    },
    # week
    {
        'level': 'week',
        'history': ('1mo', '1d'),
        'shift': 7,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + "{{{}}}",
            'polybar': _FORMAT_POLYBAR_VAL + "{{%{{F{}}}{}%{{F-}}}}%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + "{{%{{F{}}}{}%{{F-}}}}%{{B-}}%{{A}}",
        },
    },
    # month
    {
        'level': 'month',
        'history': ('3mo', '1d'),
        'shift': 30,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + "({})",
            'polybar': _FORMAT_POLYBAR_VAL + "(%{{F{}}}{}%{{F-}})%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + "(%{{F{}}}{}%{{F-}})%{{B-}}%{{A}}",
        },
    },
    # year
    {
        'level': 'year',
        'history': ('1y', '1d'),
        'shift': 365,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + "<{}>",
            'polybar': _FORMAT_POLYBAR_VAL + "<%{{F{}}}{}%{{F-}}>%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + "<%{{F{}}}{}%{{F-}}>%{{B-}}%{{A}}",
        },
    },
    # 5 years
    {
        'level': '5-years',
        'history': ('5y', '3mo'),
        'shift': 20,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + "|{}|",
            'polybar': _FORMAT_POLYBAR_VAL + "|%{{F{}}}{}%{{F-}}|%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + "|%{{F{}}}{}%{{F-}}|%{{B-}}%{{A}}",
        },
    },
    # price
    {
        'level': 'price',
        'history': ('5d', '1d'),
        'shift': 1,
        'format': {
            'plain': _FORMAT_PLAIN_VAL + ">{}<",
            'polybar': _FORMAT_POLYBAR_VAL + ">%{{F{}}}{}%{{F-}}<%{{B-}}",
            'polybar_url': _FORMAT_POLYBAR_URL_VAL + ">%{{F{}}}{}%{{F-}}<%{{B-}}%{{A}}",
        },
    },
)

_VAL_ERROR = "ERROR"
_DIFF_NA = "N/A"


class ConfigError(Exception):
    """ Configuration error class """

    def __init__(self, message):
        sys.stderr.write("%s: configuration error: %s\n" % (__program__, message))

        sys.exit(2)


class Config(object):

    def __init__(self, d):
        for k, v in d.items():
            if isinstance(v, dict):
                v = Config(v)

            setattr(self, k, v)

    def __setattr__(self, k, v):
        if not hasattr(self, k):
            super().__setattr__(k, v)
        else:
            raise RuntimeError("Can't modify read-only config setting: {}".format(k))


class ConfigFile(object):

    def __init__(self, f):
        self.config = ConfigParser()
        self.settings = argparse.Namespace(tickers=[])

        try:
            c = self.config.read(f)
        except Exception as e:
            raise ConfigError(e)
        else:
            if c:
                self.parse()
            elif id(f) != id(_DEFAULT_CONFIG):
                raise ConfigError("unable to open configuration file: {}".format(f))

    def parse(self):
        for section in _CONFIG_FORMAT:
            if section in ('general', 'colors'):
                if self.config.has_section(section):
                    try:
                        [setattr(self.settings, o, self.get(section, o, _CONFIG_FORMAT[section][o])) for o in self.config.options(section)]
                    except KeyError as e:
                        raise ConfigError("unsupported setting: {}".format(e))
            elif section == 'tickers':
                if self.config.has_section(section):
                    setattr(self.settings, section, [':'.join((o.upper(), self.config.get(section, o))) for o in self.config.options(section)])

    def get(self, s, o, t=str):
        type_map = {'str': '', 'float': 'float', 'int': 'int', 'bool': 'boolean'}

        try:
            return getattr(self.config, 'get' + type_map[t.__name__])(s, o)
        except ValueError:
            raise ConfigError("invalid setting value ({}): {}".format(t.__name__, o))


class Ticker(object):

    def __new__(cls, desc, *args, **kwargs):
        info = cls.parse(desc)

        if info[0] == 'TOTAL':
            ticker_cls = TickerTotal
        else:
            ticker_cls = TickerYahoo

        return ticker_cls(info, *args, **kwargs)

    @staticmethod
    def parse(desc):
        ti = desc.split(':', 3)

        ticker = ti[0].split('.', 1)[0]

        if len(ti) == 1:
            legend = ticker
            price = shares = 0
        elif len(ti) == 2:
            try:
                legend, price = ticker, float(ti[1])
            except ValueError:
                legend, price = ti[1], 0

            shares = 0
        elif len(ti) == 3:
            try:
                legend, price, shares = ticker, float(ti[1]), ti[2]
            except ValueError:
                legend, price, shares = ti[1], ti[2], 0
        elif len(ti) == 4:
            ticker, legend, price, shares = ti

        # empty ticker
        if not ticker:
            raise ConfigError("Ticker value can not be empty")

        # empty legend
        if not legend:
            legend = ticker

        try:
            price = float(price)
        except ValueError:
            price = 0

        try:
            shares = float(shares)
        except ValueError:
            shares = 0

        return ticker, legend, price, shares


class TickerGeneric(object):

    quote_url = None

    lock = Lock()
    _zoom = Config(_ZOOM_LEVELS[0])

    def __init__(self, info):
        self._starting = False

        self.ticker, self.legend, self.price, self.shares = info
        self.val = None
        self.diff = None

    @property
    def starting(self):
        return self._starting

    @starting.setter
    def starting(self, starting):
        if not isinstance(starting, bool):
            raise TypeError("Invalid starting value")

        self._starting = starting

    @property
    def zoom(self):
        return type(self)._zoom

    @zoom.setter
    def zoom(self, level):
        with type(self).lock:
            type(self)._zoom = level


class TickerYahoo(TickerGeneric):

    quote_url = "https\\://finance.yahoo.com/quote/"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.data = yfinance.Ticker(self.ticker)

    def fetch(self, percent=False):
        try:
            exchange = self.data.info['exchange']
            period, interval = self.zoom.history

            # account for volatility
            if self.zoom.level == 'day' and exchange == 'CCC':
                interval = "5m"

            history = self.data.history(period=period, interval=interval)

            if history.empty is True:
                raise ValueError("Historical data not found")

            val = history['Close'].iloc[-1]

            if self.zoom.level == 'price':
                if self.price > 0:
                    diff = val - self.price

                    if percent:
                        diff = diff / self.price * 100
                    elif self.shares > 0:
                        diff *= self.shares
                else:
                    diff = _DIFF_NA
            else:
                periods = self.zoom.shift

                # subtract non-trading days
                if self.zoom.level == 'year' and exchange != 'CCC':
                    periods -= 116
                # look back 24h
                elif self.zoom.level == 'day' and exchange == 'CCC':
                    periods = 288

                diff = getattr(history['Close'], 'pct_change' if percent else 'diff')(periods=periods).iloc[-2:].iloc[-1]

                if str(diff) == 'nan':
                    diff = _DIFF_NA
                elif percent:
                    diff *= 100
        except Exception:
            val, diff = _VAL_ERROR, _DIFF_NA

        self.val, self.diff = val, diff


class TickerTotal(TickerGeneric):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._tickers = []
        self.queue = Queue()
        self.threads = _DEFAULT_THREADS

    @property
    def tickers(self):
        return self._tickers

    @tickers.setter
    def tickers(self, tickers):
        self._tickers = []

        for t in tickers:
            if not isinstance(t, TickerGeneric):
                raise TypeError("Invalid ticker object provided")
            elif not isinstance(t, TickerTotal):
                self._tickers.append(t)

    def fetch(self, percent=False):
        for t in self.tickers:
            self.queue.put(t)

        threads = [Fetcher(self.queue) for _ in range(self.threads)]

        for f in threads:
            f.daemon = True
            f.start()

        self.queue.join()

        res = ([], [])

        for t in self.tickers:
            res[0].append(t.val * (t.shares if t.shares > 0 else 1) if not isinstance(t.val, str) else 0)
            res[1].append(t.diff if not isinstance(t.diff, str) else 0)

        self.val = sum(res[0])

        if self.val == 0:
            self.diff = 0
        else:
            diff_sum = sum(res[1])

            if percent:
                self.diff = diff_sum / (self.val + abs(diff_sum)) * 100
            else:
                self.diff = diff_sum


class Fetcher(Thread):

    def __init__(self, queue):
        super().__init__()

        self.queue = queue

    def run(self):
        while True:
            try:
                ticker = self.queue.get(timeout=0)

                ticker.fetch(percent=False)
            except Empty:
                break
            else:
                self.queue.task_done()


class Checker(object):

    def __init__(self, tickers, starting_ticker, percent=False, rounding=_DEFAULT_ROUNDING, padding=_DEFAULT_PADDING, thousand_separator=False, zoom_level=0, snapshot=False, conceal=False):
        self.percent = percent
        self.rounding = rounding
        self.padding = padding
        self.snapshot = snapshot
        self.conceal = conceal
        self.k = ',' if thousand_separator else ''
        self.opener = _DEFAULT_OPENER
        self.threads = _DEFAULT_THREADS
        self.colors = None
        self.current = None
        self.tickers = None
        self.zoom_levels = None
        self.exit = None
        self.skip = False

        self.init(tickers, starting_ticker, zoom_level)

        signal.signal(signal.SIGUSR1, self.cycle)
        signal.siginterrupt(signal.SIGUSR1, False)
        signal.signal(signal.SIGUSR2, self.zoom)
        signal.siginterrupt(signal.SIGUSR2, False)

    @property
    def rounded(self):
        return '{{:{k}.{:d}f}}'.format(self.rounding, k=self.k)

    @property
    def rounded_diff(self):
        return '{{:+{k}.{:d}f}}'.format(self.rounding, k=self.k)

    def concealed(self, ticker, value, diff=False):
        def conceal():
            return '{}{}.{}{}'.format(value[0] if diff else '', *((_CONCEAL_CHAR,) * 3))

        if self.conceal:
            if not diff and isinstance(ticker, TickerTotal):
                return conceal()
            elif diff and not self.percent and (isinstance(ticker, TickerTotal) or ticker.zoom.level == 'price'):
                return conceal()

        return value

    def padded(self, s):
        padding = _BASE_WIDTH + self.padding - sum(map(len, list(map(str, s))))

        return padding if padding > 0 else 0

    def init(self, tickers, starting_ticker, zoom_level):
        ticker_zero = Ticker(starting_ticker)
        ticker_list = []

        for t in tickers:
            ticker = Ticker(t)

            if ticker.ticker == ticker_zero.ticker:
                ticker.starting = True

            ticker_list.append(ticker)

        if not any(t.starting for t in ticker_list):
            raise ConfigError("{} ticker not found in the passed ticker list".format(ticker_zero.ticker))

        if self.snapshot:
            if not any(isinstance(t, TickerTotal) for t in ticker_list):
                ticker_list.append(Ticker("TOTAL"))

        for t in ticker_list:
            if isinstance(t, TickerTotal):
                if not any(t.price for t in ticker_list) and zoom_level == 'price':
                    raise ConfigError("No price per share detected for any ticker")

                t.tickers = ticker_list
                t.threads = self.threads

        self.tickers = cycle(ticker_list)
        self.zoom_levels = cycle(Config(z) for z in _ZOOM_LEVELS)
        self.exit = Event()

        while True:
            self.zoom()

            if TickerGeneric.zoom.level == zoom_level:
                break

        while True:
            self.cycle()

            if self.snapshot:
                if isinstance(self.current, TickerTotal):
                    break
            elif self.current.starting:
                break

    def decorate(self, bg, fg, hl, ga, lo):
        colors = {
            'background': bg,
            'foreground': fg,
            'highlight': hl,
            'gain': ga,
            'loss': lo
        }

        # normalize color codes
        colors.update((k, v if v[0] == '#' else '#' + v) for k, v in colors.items())

        self.colors = Config(colors)

    def display(self, ticker=None, fetch=True):
        current = ticker or self.current

        if fetch:
            current.fetch(self.percent)

        val, diff = current.val, current.diff

        if isinstance(val, str):
            rounded = val
        else:
            rounded = self.rounded.format(val)

        if isinstance(diff, str):
            rounded_diff = diff
        else:
            rounded_diff = self.rounded_diff.format(diff)

            if self.percent:
                rounded_diff += '%'

        if self.colors is None:
            output = TickerGeneric.zoom.format.plain.format(current.legend,
                                                            self.concealed(current, rounded),
                                                            self.concealed(current, rounded_diff, True))
        else:
            if isinstance(diff, str):
                diff_color = self.colors.foreground
            else:
                diff_color = self.colors.gain if diff > 0 else self.colors.loss

            output_args = [
                self.colors.highlight,
                self.colors.foreground,
                self.colors.background,
                self.colors.background,
                self.current.legend,
                self.colors.background,
                self.concealed(current, rounded),
                diff_color,
                self.concealed(current, rounded_diff, True),
            ]

            if current.quote_url is None:
                output = TickerGeneric.zoom.format.polybar.format(*output_args)
            else:
                output_args = [self.opener, current.quote_url + current.ticker] + output_args

                output = TickerGeneric.zoom.format.polybar_url.format(*output_args)

        sys.stdout.write("{}{:<{padding}}\n".format(output, "", padding=self.padded((current.legend,
                                                                                     self.concealed(current, rounded),
                                                                                     self.concealed(current, rounded_diff[1:], True)))))
        sys.stdout.flush()

    def snap(self):
        self.display()

        sys.stdout.write("===\n")
        sys.stdout.flush()

        for t in self.current.tickers:
            if t.price > 0 and self.percent and not isinstance(t.diff, str):
                t.diff = t.diff / t.shares / t.price * 100

            self.display(t, fetch=False)

    def run(self, delay=_DEFAULT_DELAY, rotate=False, one_shot=False):
        if self.snapshot:
            self.snap()
        elif one_shot:
            self.display()
        else:
            while True:
                self.exit.clear()

                while not self.exit.is_set():
                    self.display()
                    self.exit.wait(delay)

                    if rotate and not self.skip:
                        self.cycle()

                    self.skip = False

    def cycle(self, *args):
        self.skip = bool(len(args))
        self.current = next(self.tickers)
        self.exit.set()

    def zoom(self, *args):
        self.skip = bool(len(args))
        TickerGeneric.zoom = next(self.zoom_levels)
        self.exit.set()


def check_threads(n):
    try:
        i = int(n)

        if not 0 < i <= 128:
            raise ValueError
    except ValueError:
        raise argparse.ArgumentTypeError("invalid thread number: {}".format(n))

    return i


def main():
    if not HAS_YFINANCE:
        raise ConfigError("yfinance (https://github.com/ranaroussi/yfinance) is mandatory for the run")

    parser = argparse.ArgumentParser(description='Stock position / portfolio tracker for status bars')

    parser.add_argument("-c", "--percent", action='store_true',
                        help="percent change display")
    parser.add_argument("-C", "--conceal", action='store_true',
                        help="conceal calculated value")
    parser.add_argument("-d", "--delay", type=float,
                        default=_DEFAULT_DELAY,
                        help="delay between checks")
    parser.add_argument("-f", "--config-file", type=str,
                        default=_DEFAULT_CONFIG,
                        help="configuration file")
    parser.add_argument("-bg", "--background", type=str,
                        default=_DEFAULT_COLOR['background'],
                        help="background color")
    parser.add_argument("-fg", "--foreground", type=str,
                        default=_DEFAULT_COLOR['foreground'],
                        help="foreground color")
    parser.add_argument("-hl", "--highlight", type=str,
                        default=_DEFAULT_COLOR['highlight'],
                        help="ticker highlight color")
    parser.add_argument("-ga", "--gain", type=str,
                        default=_DEFAULT_COLOR['gain'],
                        help="gain accent color")
    parser.add_argument("-lo", "--loss", type=str,
                        default=_DEFAULT_COLOR['loss'],
                        help="loss accent color")
    parser.add_argument("-k", "--thousand-separator", action='store_true',
                        help="use thousand separator")
    parser.add_argument("-o", "--opener", type=str,
                        default=_DEFAULT_OPENER,
                        help="quote URL opener")
    parser.add_argument("-p", "--plain", action='store_true',
                        help="disable polybar formatting")
    parser.add_argument("-P", "--padding", type=int,
                        default=_DEFAULT_PADDING,
                        help="padding digit number")
    parser.add_argument("-r", "--rounding", type=int, choices=range(0, 14),
                        default=_DEFAULT_ROUNDING,
                        help="decimal rounding")
    parser.add_argument("-R", "--rotate", action='store_true',
                        help="automatically rotate over tickers")
    parser.add_argument("-s", "--one-shot", action='store_true',
                        help="run once and exit (first ticker only)")
    parser.add_argument("-S", "--snapshot", action='store_true',
                        help="take a snapshot")
    parser.add_argument("-t", "--starting-ticker", type=str,
                        default='',
                        help="starting ticker")
    parser.add_argument("-T", "--threads", type=check_threads,
                        default=_DEFAULT_THREADS,
                        help="number of threads to use for TOTAL ticker")
    parser.add_argument("-v", "--version", action='version',
                        version='%(prog)s {0} by {1}'.format(__version__, __author__))
    parser.add_argument("-z", "--zoom-level", type=str, choices=[l['level'] for l in _ZOOM_LEVELS],
                        default=_ZOOM_LEVELS[0]['level'],
                        help="starting zoom level")
    parser.add_argument("ticker", type=str, nargs='*',
                        metavar='TICKER[:LEGEND][:PRICE-PER-SHARE[:SHARES]]',
                        help="ticker list")

    args = parser.parse_args()

    config = ConfigFile(args.config_file)

    # merge config settings
    args = parser.parse_args(namespace=config.settings)

    tickers = args.ticker or args.tickers

    if not tickers:
        raise ConfigError("no tickers specified")

    if not args.starting_ticker:
        args.starting_ticker = tickers[0]

    checker = Checker(tickers, args.starting_ticker,
                               args.percent,
                               args.rounding,
                               args.padding,
                               args.thousand_separator,
                               args.zoom_level,
                               args.snapshot,
                               args.conceal)
    checker.threads = args.threads

    if not args.plain:
        if args.snapshot:
            raise ConfigError("snapshot can be taken only in plaintext mode")

        checker.decorate(args.background,
                         args.foreground,
                         args.highlight,
                         args.gain,
                         args.loss)
        checker.opener = args.opener

    try:
        checker.run(args.delay, args.rotate, args.one_shot)
    except KeyboardInterrupt:
        sys.stderr.write("aborted by the operator\n")

        sys.exit(2)


if __name__ == "__main__":
    main()