utils/plotting.py

import base64
import io
import matplotlib.pyplot as plt
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed
import FinanceDataReader as fdr

def get_stock_prices(stock_code, days):
    try:
        df = fdr.DataReader(stock_code)
        df = df[df.index >= df.index.max() - pd.DateOffset(days=days)]  # 최근 days일 데이터로 제한
        return df['Close']
    except Exception as e:
        print(f"Failed to fetch data for {stock_code}: {e}")
        return None

def plot_stock_prices(stock_codes, days):
    # 주식 그래프 생성을 위한 병렬 처리
    stock_prices = {}
    with ThreadPoolExecutor() as executor:
        futures = {executor.submit(get_stock_prices, stock_code.strip(), int(days)): stock_code.strip() for stock_code in stock_codes.split(',')}
        for future in as_completed(futures):
            stock_code = futures[future]
            try:
                prices = future.result()
                if prices is not None:
                    stock_prices[stock_code] = prices
            except Exception as e:
                print(f"Failed to fetch data for {stock_code}: {e}")

    # 각 주식에 대한 그래프를 그림
    plt.figure(figsize=(10, 6))
    for stock_code, prices in stock_prices.items():
        relative_prices = prices / prices.iloc[0]  # 첫 번째 데이터 포인트를 기준으로 상대적 가격 계산
        plt.plot(prices.index, relative_prices, label=stock_code.upper())  # 주식 코드를 대문자로 표시
    plt.xlabel('Date')
    plt.ylabel('Relative Price (Normalized to 1)')
    plt.title(f'Relative Stock Prices Over the Last {days} Days')
    plt.legend()

    # 그래프를 HTML로 변환하여 반환
    html_graph = io.BytesIO()
    plt.savefig(html_graph, format='png', dpi=300)
    html_graph.seek(0)
    graph_encoded = base64.b64encode(html_graph.getvalue()).decode()
    graph_html = f'<img src="data:image/png;base64,{graph_encoded}"/>'

    return graph_html