discord-stock-bot/chart_generator.py

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime
from typing import Dict, List, Optional
import logging
import io

logger = logging.getLogger(__name__)


class ChartGenerator:
    """Generates stock price charts using matplotlib."""

    @staticmethod
    def create_price_chart(
        ticker: str,
        candle_data: Dict[str, List],
        company_name: Optional[str] = None
    ) -> Optional[io.BytesIO]:
        """
        Create a price chart from candle data.

        Args:
            ticker: Stock ticker symbol
            candle_data: Dictionary with timestamps, open, high, low, close, volume
            company_name: Optional company name for title

        Returns:
            BytesIO buffer containing the chart image, or None if error
        """
        try:
            # Convert timestamps to datetime objects
            dates = [datetime.fromtimestamp(ts) for ts in candle_data['timestamps']]
            closes = candle_data['close']

            # Create figure and axis
            fig, ax = plt.subplots(figsize=(12, 6), facecolor='#1e1e1e')
            ax.set_facecolor('#2d2d2d')

            # Use index positions to remove weekend gaps
            x_pos = list(range(len(closes)))

            # Plot closing prices
            ax.plot(x_pos, closes, color='#00d4ff', linewidth=2)

            # Fill area under the line
            ax.fill_between(x_pos, closes, alpha=0.3, color='#00d4ff')

            # Format x-axis with date labels (removing gaps)
            if len(dates) > 0:
                # Check if all data is from the same day
                first_date = dates[0].date()
                last_date = dates[-1].date()
                is_intraday = first_date == last_date

                # Select tick positions (show ~7-10 labels)
                tick_spacing = max(1, len(dates) // 7)
                tick_indices = list(range(0, len(dates), tick_spacing))
                if tick_indices[-1] != len(dates) - 1:
                    tick_indices.append(len(dates) - 1)

                if is_intraday:
                    # Show time for intraday data
                    tick_labels = [dates[i].strftime('%H:%M') for i in tick_indices]
                else:
                    # Show date for multi-day data
                    tick_labels = [dates[i].strftime('%m/%d') for i in tick_indices]

                ax.set_xticks(tick_indices)
                ax.set_xticklabels(tick_labels, rotation=45, ha='right')

            # Grid styling
            ax.grid(True, alpha=0.2, color='#ffffff', linestyle='-', linewidth=0.5)

            # Labels and title
            title = f"{company_name} ({ticker})" if company_name else ticker
            ax.set_title(title, color='#ffffff', fontsize=14, fontweight='bold', pad=20)
            ax.set_xlabel('Date', color='#ffffff', fontsize=10)
            ax.set_ylabel('Price ($)', color='#ffffff', fontsize=10)

            # Tick colors
            ax.tick_params(colors='#ffffff', which='both')

            # Spine colors
            for spine in ax.spines.values():
                spine.set_color('#444444')

            # Tight layout
            plt.tight_layout()

            # Save to buffer
            buffer = io.BytesIO()
            plt.savefig(buffer, format='png', facecolor='#1e1e1e', dpi=100)
            buffer.seek(0)
            plt.close(fig)

            return buffer

        except Exception as e:
            logger.error(f"Error generating chart for {ticker}: {e}")
            return None

    @staticmethod
    def create_candlestick_chart(
        ticker: str,
        candle_data: Dict[str, List],
        company_name: Optional[str] = None
    ) -> Optional[io.BytesIO]:
        """
        Create a candlestick chart from candle data.

        Args:
            ticker: Stock ticker symbol
            candle_data: Dictionary with timestamps, open, high, low, close, volume
            company_name: Optional company name for title

        Returns:
            BytesIO buffer containing the chart image, or None if error
        """
        try:
            # Convert timestamps to datetime objects
            dates = [datetime.fromtimestamp(ts) for ts in candle_data['timestamps']]
            opens = candle_data['open']
            highs = candle_data['high']
            lows = candle_data['low']
            closes = candle_data['close']

            # Create figure and axis
            fig, ax = plt.subplots(figsize=(12, 6), facecolor='#1e1e1e')
            ax.set_facecolor('#2d2d2d')

            # Calculate bar width
            width = 0.6

            # Plot candlesticks using index positions (removes weekend gaps)
            for i in range(len(dates)):
                color = '#00ff88' if closes[i] >= opens[i] else '#ff4444'

                # High-low line
                ax.plot([i, i], [lows[i], highs[i]], color=color, linewidth=1)

                # Open-close rectangle
                height = abs(closes[i] - opens[i])
                bottom = min(opens[i], closes[i])
                ax.bar(i, height, width, bottom=bottom, color=color, alpha=0.8)

            # Format x-axis with date labels (removing gaps)
            if len(dates) > 0:
                # Check if all data is from the same day
                first_date = dates[0].date()
                last_date = dates[-1].date()
                is_intraday = first_date == last_date

                # Select tick positions (show ~7-10 labels)
                tick_spacing = max(1, len(dates) // 7)
                tick_indices = list(range(0, len(dates), tick_spacing))
                if tick_indices[-1] != len(dates) - 1:
                    tick_indices.append(len(dates) - 1)

                if is_intraday:
                    # Show time for intraday data
                    tick_labels = [dates[i].strftime('%H:%M') for i in tick_indices]
                else:
                    # Show date for multi-day data
                    tick_labels = [dates[i].strftime('%m/%d') for i in tick_indices]

                ax.set_xticks(tick_indices)
                ax.set_xticklabels(tick_labels, rotation=45, ha='right')

            # Grid styling
            ax.grid(True, alpha=0.2, color='#ffffff', linestyle='-', linewidth=0.5)

            # Labels and title
            title = f"{company_name} ({ticker})" if company_name else ticker
            ax.set_title(title, color='#ffffff', fontsize=14, fontweight='bold', pad=20)
            ax.set_xlabel('Date', color='#ffffff', fontsize=10)
            ax.set_ylabel('Price ($)', color='#ffffff', fontsize=10)

            # Tick colors
            ax.tick_params(colors='#ffffff', which='both')

            # Spine colors
            for spine in ax.spines.values():
                spine.set_color('#444444')

            # Tight layout
            plt.tight_layout()

            # Save to buffer
            buffer = io.BytesIO()
            plt.savefig(buffer, format='png', facecolor='#1e1e1e', dpi=100)
            buffer.seek(0)
            plt.close(fig)

            return buffer

        except Exception as e:
            logger.error(f"Error generating candlestick chart for {ticker}: {e}")
            return None