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

@@ -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

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

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)

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

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 = (
+        "<br> <br> <b>%{label}<br>    %{customdata[0]:.2f}% <br> <br> <br><br><b>"
+    )
+    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()

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
 

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.
+```

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: