diff --git a/openbb_terminal/stocks/comparison_analysis/finbrain_model.py b/openbb_terminal/stocks/comparison_analysis/finbrain_model.py index 293bdd64700..fe5a69ea9f4 100644 --- a/openbb_terminal/stocks/comparison_analysis/finbrain_model.py +++ b/openbb_terminal/stocks/comparison_analysis/finbrain_model.py @@ -13,6 +13,36 @@ from openbb_terminal.rich_config import console logger = logging.getLogger(__name__) +@log_start_end(log=logger) +def find_smallest_num_data_point(results_list: List[dict]) -> int: + """Helper function to find the number for the smallest total number of data points + out of all the tickers. This is necessary because if one ticker has more or less + data points than another then it will throw an indexing error. The solution is to + have each ticker have the same number of data points as to graph and view properly. + We chose to set each ticker to the minimum of number data points out of all the + tickers + + Parameters + ---------- + results_list : List[json] + List of dicts storing ticker data + + Returns + ------- + int + Value of smallest total number of sentiment data points + """ + small_list = list() + for result_json in results_list: + if ( + "ticker" in result_json + and "sentimentAnalysis" in result_json + and len(result_json["sentimentAnalysis"].values()) > 0 + ): + small_list.append(len(result_json["sentimentAnalysis"].values())) + return min(small_list) + + @log_start_end(log=logger) def get_sentiments(tickers: List[str]) -> pd.DataFrame: """Gets Sentiment analysis from several tickers provided by FinBrain's API @@ -31,26 +61,54 @@ def get_sentiments(tickers: List[str]) -> pd.DataFrame: df_sentiment = pd.DataFrame() dates_sentiment = [] tickers_to_remove = list() + results_list = list() for ticker in tickers: result = requests.get(f"https://api.finbrain.tech/v0/sentiments/{ticker}") + # Check status code, if its correct then convert to dict using .json() if result.status_code == 200: result_json = result.json() - if "ticker" in result_json and "sentimentAnalysis" in result_json: - df_sentiment[ticker] = [ - float(val) - for val in list(result_json["sentimentAnalysis"].values()) - ] - dates_sentiment = list(result_json["sentimentAnalysis"].keys()) - else: - console.print(f"Unexpected data format from FinBrain API for {ticker}") - tickers_to_remove.append(ticker) - + results_list.append(result_json) else: console.print( f"Request error in retrieving {ticker} sentiment from FinBrain API" ) tickers_to_remove.append(ticker) + # Finds the smallest amount of data points from any of the tickers as to not run + # into an indexing error when graphing + smallest_num_data_point = find_smallest_num_data_point(results_list) + + num = 0 + for result_json in results_list: + ticker = tickers[num] + # Checks to see if sentiment data in results_json + if ( + "ticker" in result_json + and "sentimentAnalysis" in result_json + and len(result_json["sentimentAnalysis"].values()) > 0 + ): + # Storing sentiments and dates in list + sentiments = list(result_json["sentimentAnalysis"].values()) + dates_sentiment = list(result_json["sentimentAnalysis"].keys()) + + # If there are more sentiment data points for one ticker compared to the + # smallest amount of data points, then remove that data points from that + # ticker as to be able to graph properly + if len(sentiments) > smallest_num_data_point: + sentiments = sentiments[0:smallest_num_data_point] + dates_sentiment = dates_sentiment[0:smallest_num_data_point] + + df_sentiment[ticker] = [float(val) for val in sentiments] + + # If sentiment data not in results_json remove it + else: + console.print( + f"Unexpected data format or no data from FinBrain API for {ticker}" + ) + tickers_to_remove.append(ticker) + + num = num + 1 + for ticker in tickers_to_remove: tickers.remove(ticker)