# -*- coding: utf-8 -*- from functools import lru_cache from pathlib import Path from random import random import pandas as pd from alphaVantageAPI.alphavantage import AlphaVantage # pip install alphaVantage-api import pandas_ta as ta class Watchlist(object): """Watchlist Class (** This is subject to change! **) ============================================================================ A simple Class to load/download financial market data and automatically apply Technical Analysis indicators with a Pandas TA Strategy. Default Strategy: pandas_ta.AllStrategy. Requirements: - Pandas TA (pip install pandas_ta) - AlphaVantage (pip install alphaVantage-api) for the Default Data Source. To use another Data Source, update the load() method after AV. Required Arguments: - tickers: A list of strings containing tickers. Example: ['SPY', 'AAPL'] ============================================================================ """ def __init__( self, tickers: list, tf: str = None, name: str = None, strategy: ta.Strategy = None, ds: object = None, **kwargs ): self.tickers = tickers self.tf = tf self.verbose = kwargs.pop("verbose", False) self.name = name self.data = None self.kwargs = kwargs self.ds = ds if ds is not None else None self.strategy = strategy def _drop_columns(self, df: pd.DataFrame, cols: list = ['Unnamed: 0', 'date', 'split_coefficient', 'dividend']): """Helper methods to drop columns silently.""" df_columns = list(df.columns) if any(_ in df_columns for _ in cols): if self.verbose: print(f"[i] Possible columns dropped: {', '.join(cols)}") df = df.drop(cols, axis=1, errors='ignore') return df def _load_all(self, **kwargs) -> dict: """Updates the Watchlist's data property with a dictionary of DataFrames keyed by ticker.""" if self.tickers is not None and isinstance(self.tickers, list) and len(self.tickers): self.data = {ticker: self.load(ticker, **kwargs) for ticker in self.tickers} return self.data def load( self, ticker: str = None, tf: str = None, index: str = 'date', drop: list = ['dividend', 'split_coefficient'], file_path: str = ".", **kwargs ) -> pd.DataFrame: """Loads or Downloads (if a local csv does not exist) the data from the Data Source. When successful, it returns a Data Frame for the requested ticker. If no tickers are given, it loads all the tickers.""" tf = self.tf if tf is None else tf.upper() if ticker is not None and isinstance(ticker, str): ticker = str(ticker).upper() else: print(f"[!] Loading All: {', '.join(self.tickers)}") self._load_all(**kwargs) return filename_ = f"{ticker}_{tf}.csv" current_file = Path(file_path) / filename_ # Load local or from Data Source if current_file.exists(): df = pd.read_csv(filename_, index_col=index) if not df.ta.datetime_ordered: df = df.set_index(pd.DatetimeIndex(df.index)) print(f"\n[i] Loaded['{tf}']: {filename_}") else: if self.ds is not None and isinstance(self.ds, AlphaVantage): df = self.ds.data(tf, ticker) if not df.ta.datetime_ordered: df = df.set_index(pd.DatetimeIndex(df[index])) print(f"\n[+] Downloading['{tf}']: {ticker}") df = self._drop_columns(df) # Remove select columns if kwargs.pop("analyze", True): df.ta.strategy(name=self.strategy.name, ta=self.strategy.ta, **kwargs) df.ticker = ticker # Attach ticker to the DataFrame return df @property def data(self) -> dict: """When not None, it contains a dictionary of DataFrames keyed by ticker. data = {"SPY": pd.DataFrame, ...}""" return self._data @data.setter def data(self, value: dict) -> None: # Later check dict has string keys and DataFrame values if value is not None and isinstance(value, dict): if self.verbose: print(f"[+] New data") self._data = value else: self._data = None @property def name(self) -> str: """The name of the Watchlist. Default: "Watchlist: {Watchlist.tickers}".""" return self._name @name.setter def name(self, value: str) -> None: if isinstance(value, str): self._name = str(value) else: self._name = f"Watchlist: {', '.join(self.tickers)}" @property def strategy(self) -> ta.Strategy: """Pandas TA Strategy Class. Default: pandas_ta.AllStrategy""" return self._strategy @strategy.setter def strategy(self, value: ta.Strategy) -> None: if value is not None and isinstance(value, ta.Strategy): self._strategy = value else: self._strategy = ta.AllStrategy @property def tf(self) -> str: """Alias for timeframe. Default: 'D'""" return self._tf @tf.setter def tf(self, value: str) -> None: if isinstance(value, str): value = str(value) self._tf = value else: self._tf = "D" @property def tickers(self) -> list: """tickers If a string, it it converted to a list. Example: 'AAPL' -> ['AAPL'] * Does not accept, comma seperated strings. If a list, checks if it is a list of strings. """ return self._tickers @tickers.setter def tickers(self, value: (list, str)) -> None: if value is None: print(f"[X] {value} is not a valie Watchlist ticker.") return elif isinstance(value, list) and [isinstance(_, str) for _ in value]: self._tickers = list(map(str.upper, value)) elif isinstance(value, str): self._tickers = [value.upper()] self.name = self._tickers @property def verbose(self) -> bool: """Toggle the verbose property. Default: False""" return self._verbose @verbose.setter def verbose(self, value: bool) -> None: if isinstance(value, bool): self._verbose = bool(value) else: self._verbose = False def indicators(self, *args, **kwargs) -> any: """Returns the list of indicators that are available with Pandas Ta.""" pd.DataFrame().ta.indicators(*args, **kwargs) def __repr__(self) -> str: s = f"Watch(name='{self.name}', tickers[{len(self.tickers)}]='{', '.join(self.tickers)}', tf='{self.tf}', strategy[{self.strategy.total_ta()}]='{self.strategy.name}'" if self.data is not None: s += f", data[{len(self.data.keys())}])" return s return s + ")"