From f856a2f2125e58839f21e74f3dd84dfe598d515e Mon Sep 17 00:00:00 2001 From: James Maslek Date: Tue, 6 Sep 2022 21:15:47 -0400 Subject: [PATCH] Add heatmap function to stocks disc (#2454) * Add heatmap function * Reinstalled black. Weird * fix export and lint * fix test Co-authored-by: Colin Delahunty <72827203+colin99d@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- i18n/en.yml | 1 + .../stocks/discovery/disc_controller.py | 69 +++++++---- .../stocks/discovery/finviz_model.py | 58 +++++++++ .../stocks/discovery/finviz_view.py | 111 ++++++++++++++++++ .../test_disc_controller/test_print_help.txt | 1 + .../terminal/stocks/disc/heatmap/_index.md | 15 +++ website/data/menu/main.yml | 4 +- 8 files changed, 238 insertions(+), 23 deletions(-) create mode 100644 openbb_terminal/stocks/discovery/finviz_model.py create mode 100644 openbb_terminal/stocks/discovery/finviz_view.py create mode 100644 website/content/terminal/stocks/disc/heatmap/_index.md diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f1d3695683c..807d8fe8612 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -32,7 +32,7 @@ repos: - id: flake8 entry: flake8 --ignore=E203,W503 --max-line-length=122 --exclude ./build/pyinstaller/* - repo: https://github.com/codespell-project/codespell - rev: v2.2.0 + rev: v2.2.1 hooks: - id: codespell entry: codespell diff --git a/i18n/en.yml b/i18n/en.yml index 8c183457cf5..5cadee94f60 100644 --- a/i18n/en.yml +++ b/i18n/en.yml @@ -191,6 +191,7 @@ en: stocks/disc/hotpenny: today's hot penny stocks stocks/disc/rtat: top 10 retail traded stocks per day stocks/disc/divcal: dividend calendar for selected date + stocks/disc/heatmap: heatmap of SP500 tickers stocks/sia/load: load a specific ticker and all it's corresponding parameters stocks/sia/clear: clear all or one of industry, sector, country and market cap parameters stocks/sia/industry: see existing industries, or set industry if arg specified diff --git a/openbb_terminal/stocks/discovery/disc_controller.py b/openbb_terminal/stocks/discovery/disc_controller.py index e46b43a313a..ca803c99891 100644 --- a/openbb_terminal/stocks/discovery/disc_controller.py +++ b/openbb_terminal/stocks/discovery/disc_controller.py @@ -29,6 +29,7 @@ from openbb_terminal.stocks.discovery import ( seeking_alpha_view, shortinterest_view, yahoofinance_view, + finviz_view, ) # pylint:disable=C0302 @@ -59,6 +60,7 @@ class DiscoveryController(BaseController): "cnews", "rtat", "divcal", + "heatmap", ] arkord_sortby_choices = [ @@ -110,6 +112,7 @@ class DiscoveryController(BaseController): "Annual Dividend", "Announcement Date", ] + heatmap_timeframes = ["day", "week", "month", "3month", "6month", "year", "ytd"] def __init__(self, queue: List[str] = None): """Constructor""" @@ -126,30 +129,31 @@ class DiscoveryController(BaseController): choices["cnews"]["--type"] = {c: None for c in self.cnews_type_choices} choices["divcal"]["-s"] = {c: None for c in self.dividend_columns} choices["divcal"]["--sort"] = {c: None for c in self.dividend_columns} - + choices["heatmap"]["-t"] = {c: None for c in self.heatmap_timeframes} self.completer = NestedCompleter.from_nested_dict(choices) def print_help(self): """Print help""" mt = MenuText("stocks/disc/") - mt.add_cmd("pipo") - mt.add_cmd("fipo") - mt.add_cmd("gainers") - mt.add_cmd("losers") - mt.add_cmd("ugs") - mt.add_cmd("gtech") - mt.add_cmd("active") - mt.add_cmd("ulc") - mt.add_cmd("asc") - mt.add_cmd("ford") - mt.add_cmd("arkord") - mt.add_cmd("upcoming") - mt.add_cmd("trending") - mt.add_cmd("cnews") - mt.add_cmd("lowfloat") - mt.add_cmd("hotpenny") - mt.add_cmd("rtat") - mt.add_cmd("divcal") + mt.add_cmd("pipo", "Finnhub") + mt.add_cmd("fipo", "Finnhub") + mt.add_cmd("gainers", "Yahoo Finance") + mt.add_cmd("losers", "Yahoo Finance") + mt.add_cmd("ugs", "Yahoo Finance") + mt.add_cmd("gtech", "Yahoo Finance") + mt.add_cmd("active", "Yahoo Finance") + mt.add_cmd("ulc", "Yahoo Finance") + mt.add_cmd("asc", "Yahoo Finance") + mt.add_cmd("ford", "Fidelity") + mt.add_cmd("arkord", "Cathies Ark") + mt.add_cmd("upcoming", "Seeking Alpha") + mt.add_cmd("trending", "Seeking Alpha") + mt.add_cmd("cnews", "Seeking Alpha") + mt.add_cmd("lowfloat", "Fidelity") + mt.add_cmd("hotpenny", "Shortinterest") + mt.add_cmd("rtat", "NASDAQ Data Link") + mt.add_cmd("divcal", "NASDAQ Data Link") + mt.add_cmd("heatmap", "Finviz") console.print(text=mt.menu_text, menu="Stocks - Discovery") # TODO Add flag for adding last price to the following table @@ -835,7 +839,7 @@ class DiscoveryController(BaseController): @log_start_end(log=logger) def call_rtat(self, other_args: List[str]): - """Process fds command""" + """Process rtat command""" parser = argparse.ArgumentParser( add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter, @@ -864,3 +868,28 @@ class DiscoveryController(BaseController): nasdaq_view.display_top_retail( limit=ns_parser.limit, export=ns_parser.export ) + + @log_start_end(log=logger) + def call_heatmap(self, other_args: List[str]): + """Process heatmap command""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="heatmap", + description=""" + Get the SP 500 heatmap from finviz and display in interactive treemap + """, + ) + parser.add_argument( + "-t", + "--timeframe", + default="day", + choices=self.heatmap_timeframes, + help="Timeframe to get heatmap data for", + dest="timeframe", + ) + ns_parser = self.parse_known_args_and_warn( + parser, other_args, export_allowed=EXPORT_ONLY_RAW_DATA_ALLOWED + ) + if ns_parser: + finviz_view.display_heatmap(ns_parser.timeframe, ns_parser.export) diff --git a/openbb_terminal/stocks/discovery/finviz_model.py b/openbb_terminal/stocks/discovery/finviz_model.py new file mode 100644 index 00000000000..fc3fdef6f53 --- /dev/null +++ b/openbb_terminal/stocks/discovery/finviz_model.py @@ -0,0 +1,58 @@ +"""Finviz Model""" +__docformat__ = "numpy" + +import pandas as pd +import requests +from openbb_terminal.helper_funcs import get_user_agent +from openbb_terminal.rich_config import console + +timeframe_map = { + "day": "", + "week": "w1", + "month": "w4", + "3month": "w13", + "6month": "w26", + "year": "w52", + "ytd": "ytd", +} + + +def get_heatmap_data(timeframe: str) -> pd.DataFrame: + """Get heatmap data from finviz + + Parameters + ---------- + timeframe: str + Timeframe to get performance for + Returns + ------- + pd.DataFrame + Dataframe of tickers, changes and sectors + """ + if timeframe not in timeframe_map: + console.print(f"[red]{timeframe} is an invalid timeframe[/red]") + return pd.DataFrame() + r = requests.get( + f"https://finviz.com/api/map_perf.ashx?t=sec&st={timeframe_map[timeframe]}", + headers={"User-Agent": get_user_agent()}, + ) + r2_json = requests.get( + "https://finviz.com/maps/sec.json?rev=316", + headers={"User-Agent": get_user_agent()}, + ).json() + df_change = pd.DataFrame.from_dict(r.json()["nodes"], orient="index") + df_change.columns = ["Change"] + df_change["Change"] = df_change.Change / 100 + changes_dict = df_change.to_dict() + + dfs = pd.DataFrame() + for sector_dict in r2_json["children"]: + for industry_dict in sector_dict["children"]: + temp = pd.DataFrame(industry_dict["children"]) + temp["Sector"] = sector_dict["name"] + temp["Industry"] = industry_dict["name"] + dfs = pd.concat([dfs, temp], axis=0).reset_index(drop=True) + + dfs["Change"] = dfs["name"].copy().map(changes_dict["Change"]) + dfs = dfs.rename(columns={"name": "Ticker"}) + return dfs diff --git a/openbb_terminal/stocks/discovery/finviz_view.py b/openbb_terminal/stocks/discovery/finviz_view.py new file mode 100644 index 00000000000..dd7e550a46e --- /dev/null +++ b/openbb_terminal/stocks/discovery/finviz_view.py @@ -0,0 +1,111 @@ +"""Finviz View""" +__docformat__ = "numpy" + +import os +import pandas as pd + +import plotly.express as px +from plotly.subplots import make_subplots + +from openbb_terminal.stocks.discovery import finviz_model +from openbb_terminal.rich_config import console +from openbb_terminal.helper_funcs import export_data + + +def display_heatmap(timeframe: str, export: str = ""): + """Display heatmap from finviz + + Parameters + ---------- + timeframe: str + Timeframe for performance + export: str + Format to export data + """ + console.print("") + dfs = finviz_model.get_heatmap_data(timeframe) + if dfs.empty: + return + export_data( + export, + os.path.dirname(os.path.abspath(__file__)), + "heatmap", + dfs, + ) + color_bin = [-1, -0.02, -0.01, -0.00001, 0.00001, 0.01, 0.02, 1] + dfs["colors"] = pd.cut( + dfs["Change"], + bins=color_bin, + labels=[ + "rgb(246, 53, 56)", + "rgb(191, 64, 69)", + "rgb(139, 68, 78)", + "grey", + "rgb(53, 118, 78)", + "rgb(47, 158, 79)", + "rgb(48, 204, 90)", + ], + ) + path_tree = [px.Constant("SP 500"), "Sector", "Ticker"] + fig = make_subplots( + print_grid=False, + vertical_spacing=0.02, + horizontal_spacing=-0, + specs=[[{"type": "domain"}]], + rows=1, + cols=1, + ) + treemap = px.treemap( + dfs, + path=path_tree, + values="value", + custom_data=["Change"], + color="colors", + color_discrete_map={ + "(?)": "#262931", + "grey": "grey", + "rgb(246, 53, 56)": "rgb(246, 53, 56)", + "rgb(191, 64, 69)": "rgb(191, 64, 69)", + "rgb(139, 68, 78)": "rgb(139, 68, 78)", + "rgb(53, 118, 78)": "rgb(53, 118, 78)", + "rgb(47, 158, 79)": "rgb(47, 158, 79)", + "rgb(48, 204, 90)": "rgb(48, 204, 90)", + }, + ) + fig.add_trace(treemap["data"][0], row=1, col=1) + + fig.data[ + 0 + ].texttemplate = ( + "

%{label}
%{customdata[0]:.2f}%



" + ) + fig.data[0].insidetextfont = dict( + family="Arial Black", + size=50, + color="white", + ) + + fig.update_traces( + textinfo="label+text+value", + textposition="middle center", + selector=dict(type="treemap"), + marker_line_width=0.3, + marker_pad_b=20, + marker_pad_l=0, + marker_pad_r=0, + marker_pad_t=50, + tiling_pad=2, + ) + fig.update_layout( + margin=dict(t=0, l=0, r=0, b=0), + paper_bgcolor="rgba(0,0,0,0)", + plot_bgcolor="rgba(0,0,0,0)", + hovermode=False, + font=dict( + family="Arial Black", + size=20, + color="white", + ), + ) + + fig.show() diff --git a/tests/openbb_terminal/stocks/discovery/txt/test_disc_controller/test_print_help.txt b/tests/openbb_terminal/stocks/discovery/txt/test_disc_controller/test_print_help.txt index a1d80967b07..4431e8d70c8 100644 --- a/tests/openbb_terminal/stocks/discovery/txt/test_disc_controller/test_print_help.txt +++ b/tests/openbb_terminal/stocks/discovery/txt/test_disc_controller/test_print_help.txt @@ -16,4 +16,5 @@ hotpenny today's hot penny stocks [YahooFinance, Shortinterest] rtat top 10 retail traded stocks per day [Nasdaq] divcal dividend calendar for selected date [Nasdaq] + heatmap heatmap of SP500 tickers diff --git a/website/content/terminal/stocks/disc/heatmap/_index.md b/website/content/terminal/stocks/disc/heatmap/_index.md new file mode 100644 index 00000000000..e4ee6181d11 --- /dev/null +++ b/website/content/terminal/stocks/disc/heatmap/_index.md @@ -0,0 +1,15 @@ +``` +2022 Aug 29, 13:38 (🦋) /stocks/disc/ $ heatmap -h +usage: heatmap [-t {day,week,month,3month,6month,year,ytd}] [-h] [--export EXPORT] +``` + +Get the SP 500 heatmap from finviz and display in interactive treemap. +``` +options: + -t {day,week,month,3month,6month,year,ytd}, --timeframe {day,week,month,3month,6month,year,ytd} + Timeframe to get heatmap data for (default: day) + -h, --help show this help message (default: False) + --export EXPORT Export raw data into csv, json, xlsx (default: ) + +For more information and examples, use 'about heatmap' to access the related guide. +``` \ No newline at end of file diff --git a/website/data/menu/main.yml b/website/data/menu/main.yml index b83b1861863..0359c611560 100755 --- a/website/data/menu/main.yml +++ b/website/data/menu/main.yml @@ -59,8 +59,8 @@ main: ref: "/terminal/stocks/disc/rtat" - name: divcal ref: "/terminal/stocks/disc/divcal" - - name: cramer - ref: "/terminal/stocks/disc/cramer" + - name: heatmap + ref: "/terminal/stocks/disc/heatmap" - name: Sector and Industry Analysis (SIA) ref: "/terminal/stocks/sia" sub: