예제) matplotlib으로 주식 차트 그리기(matplotlib candlestick chart example)

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.axes import Axes
from matplotlib.widgets import SpanSelector

class Chart:
    def __init__(self, df: pd.DataFrame, date='기준일', close='종가', Open='시가', high='고가', low='저가', volume='거래량', ma=[5, 10, 20, 60, 120, 240],) -> None:
        df.set_index(date)
        for i in ma: df[f'{i}일선'] = df[close].rolling(i).mean()

        self.df, self.date, self.close = (df, date, close)
        self.open, self.high, self.low, self.volume = (Open, high, low, volume)
        self.ma = ma
        self.set_fig()
        self.draw()

        return

    def set_fig(self):
        df, close = (self.df, self.close)
        Open, high, low, volume = (self.open, self.high, self.low, self.volume)
        list_ma = self.ma

        fig, ax = plt.subplots(3, figsize=(8, 4))
        ax: list[Axes]
        ax1, ax2, ax3 = ax

        ax2.sharex(ax3)

        x = np.arange(len(df.index))
        y = df[[Open, high, low, close]].values
        end_y_index = len(y) - 1
        ax1.plot(x, df[self.close].values)
        ax1.set_ylim(df[low].min() * 0.9, df[high].max() * 1.1)
        ax1.set_facecolor('gray')
        ax1.set_title('Press left mouse button and drag '
                    'to select a region in the top graph')

        red, blue = ([], [])
        for index, i in df.iterrows():
            c, o, h, l = (i[close], i[Open], i[high], i[low])
            red.append(None), blue.append(None)
            # 양봉
            if o <= c:  red[-1] = index
            # 음봉
            else:  blue[-1] = index

        # 양봉 그리기
        ax2.vlines(df.index, [df[Open][i] if i else None for i in red], [df[close][i] if i else None for i in red], color='r', linestyle='-', lw=10)
        ax2.vlines(df.index, [df[low][i] if i else None for i in red], [df[high][i] if i else None for i in red], color='r', linestyle='-', lw=1)
        # 음봉 그리기
        ax2.vlines(df.index, [df[Open][i] if i else None for i in blue], [df[close][i] if i else None for i in blue], color='b', linestyle='-', lw=10)
        ax2.vlines(df.index, [df[low][i] if i else None for i in blue], [df[high][i] if i else None for i in blue], color='b', linestyle='-', lw=1)
        ax2.set_ylim(df[low].min() * 0.9, df[high].max() * 1.1)
        for i in list_ma: ax2.plot(x, df[f'{i}일선'])
        ax3.bar(x, df[volume].values)

        self.x, self.y, self.end_y_index = (x, y, end_y_index)
        self.fig, self.ax1, self.ax2, self.ax3 = (fig, ax1, ax2, ax3)

        return

    axvspan = None
    def onselect(self, xmin, xmax):
        x, y = (self.x, self.y)
        fig, ax1, ax2, ax3 = (self.fig, self.ax1, self.ax2, self.ax3)

        indmin, base_indmax = np.searchsorted(x, (xmin, xmax))
        indmax = min(len(x)-1, base_indmax)

        region_x = x[indmin:indmax]
        region_y = y[indmin:indmax]

        if len(region_x) >= 2:
            if self.axvspan: self.axvspan.remove()
            self.axvspan = ax1.axvspan(xmin, xmax, facecolor='white', alpha=1.0)

            ax2.set_xlim(region_x[0], region_x[-1])
            ax2.set_ylim(region_y.min() * 0.9, region_y.max() * 1.1)
            region_y_volume = df[self.volume][indmin:indmax]
            ax3.set_ylim(region_y_volume.min(), region_y_volume.max() * 1.1)
            fig.canvas.draw_idle()
        return

    def draw(self):
        ax1 = self.ax1
        span = SpanSelector(
            ax1,
            self.onselect,
            'horizontal',
            useblit=True,
            props=dict(alpha=0.0, facecolor='tab:blue'),
            interactive=True,
            drag_from_anywhere=True
        )

        plt.show()

        return


if __name__ == '__main__':
    list_price = {주식 시세 데이터}
    df = pd.DataFrame(list_price)
    a = Chart(df)

from time import time

import matplotlib.pyplot as plt
import matplotlib.style as mplstyle
from matplotlib.backend_bases import MouseEvent
import pandas as pd
from matplotlib.axes import Axes
from matplotlib.widgets import SpanSelector

mplstyle.use('fast')
plt.rcParams['font.family'] ='Malgun Gothic'

class BlittedCursor:
    def __init__(self, ax: Axes, x, y, ax_sub: Axes, y_sub,):
        self.ax = ax
        self.x, self.y = (x, y)
        self.background = None
        self.horizontal_line = ax.axhline(color='k', lw=0.8, ls='-')
        self.vertical_line = ax.axvline(color='k', lw=0.8, ls='-')
        self.text_price = ax.text(0, 0, '', horizontalalignment='left', verticalalignment='bottom', bbox=dict(boxstyle='round', facecolor='w', zorder=5))
        ax.figure.canvas.mpl_connect('draw_event', self.on_draw)
        self.ax_sub = ax_sub
        self.y_sub = y_sub
        self.background_sub = None
        self.horizontal_line_sub = ax_sub.axhline(color='k', lw=0.8, ls='-')
        self.vertical_line_sub = ax_sub.axvline(color='k', lw=0.8, ls='-')
        self.text_date = ax_sub.text(0, 0, '', horizontalalignment='center', verticalalignment='top', bbox=dict(boxstyle='round', facecolor='w', zorder=5))
        self.text_volume = ax_sub.text(0, 0, '', horizontalalignment='left', verticalalignment='bottom', bbox=dict(boxstyle='round', facecolor='w', zorder=5))
        ax_sub.figure.canvas.mpl_connect('draw_event', self.on_draw)
        self._creating_background = False
        return
    def on_draw(self, event):
        self.create_new_background()
    def set_cross_hair_visible(self, visible):
        need_redraw = self.horizontal_line.get_visible() != visible
        self.horizontal_line.set_visible(visible)
        self.vertical_line.set_visible(visible)
        self.text_price.set_visible(visible)
        self.horizontal_line_sub.set_visible(visible)
        self.vertical_line_sub.set_visible(visible)
        self.text_volume.set_visible(visible)
        self.text_date.set_visible(visible)
        return need_redraw
    def create_new_background(self):
        if self._creating_background: return
        self._creating_background = True
        self.set_cross_hair_visible(False)
        self.ax.figure.canvas.draw()
        self.background = self.ax.figure.canvas.copy_from_bbox(self.ax.bbox)
        self.ax_sub.figure.canvas.draw()
        self.background_sub = self.ax_sub.figure.canvas.copy_from_bbox(self.ax_sub.bbox)
        self.set_cross_hair_visible(True)
        self._creating_background = False
        return
    def on_mouse_move(self, event: MouseEvent):
        if self.background is None: self.create_new_background()
        ax = event.inaxes
        if not ax:
            need_redraw = self.set_cross_hair_visible(False)
            if need_redraw:
                self.ax.figure.canvas.restore_region(self.background)
                self.ax.figure.canvas.blit(self.ax.bbox)
                self.ax_sub.figure.canvas.restore_region(self.background_sub)
                self.ax_sub.figure.canvas.blit(self.ax_sub.bbox)
        else:
            self.set_cross_hair_visible(True)
            x, y = event.xdata, event.ydata
            indx, indy = (int(x), int(y))
            self.ax.figure.canvas.restore_region(self.background)
            self.ax_sub.figure.canvas.restore_region(self.background_sub)
            self.vertical_line.set_xdata([x])
            self.ax.draw_artist(self.vertical_line)
            self.vertical_line_sub.set_xdata([x])
            self.ax_sub.draw_artist(self.vertical_line_sub)
            try: text_date = f'{self.x[indx]}'
            except KeyError as e: text_date = ''
            self.text_date.set_text(text_date)
            self.text_date.set_x(x)
            _, ymax = self.ax_sub.get_ylim()
            if ax is self.ax:
                self.vertical_line_sub.set_visible(False)
                self.text_date.set_y(ymax * 0.9)
                if indy < 0: price = ''
                else:
                    try: price = f'{self.y_sub[indy]:,}원'
                    except KeyError as e: price = ''
                self.text_price.set_text(price)
                self.text_price.set_position([ax.get_xlim()[0]+1, y])
                self.ax.draw_artist(self.text_price)
                self.horizontal_line.set_ydata([y])
                self.ax.draw_artist(self.horizontal_line)
            elif ax is self.ax_sub:
                self.vertical_line.set_visible(False)
                self.text_date.set_y(ymax * 1.25)
                try: volume = convert_unit(self.y_sub[indy], digit=2, word='주')
                except KeyError as e: volume = ''
                self.text_volume.set_text(volume)
                self.text_volume.set_position([self.ax_sub.get_xlim()[0]+1, y])
                self.ax_sub.draw_artist(self.text_volume)
                self.horizontal_line_sub.set_ydata([y])
                self.ax_sub.draw_artist(self.horizontal_line_sub)
            else:
                self.vertical_line_sub.set_visible(False)
                self.vertical_line.set_visible(False)
                self.text_price.set_visible(False)
                self.text_volume.set_visible(False)
            self.ax_sub.draw_artist(self.text_date)
            self.ax.figure.canvas.blit(self.ax.bbox)
            self.ax_sub.figure.canvas.blit(self.ax_sub.bbox)
            return

dict_unit = {
    '조': 1_000_000_000_000,
    '억':       100_000_000,
    '만':            10_000,
    '천':             1_000,
}
def convert_unit(value, digit=0, word='원'):
    for unit, n in dict_unit.items():
        if n <= value:
            num = round(value / n, digit)
            if not num % 1: num = int(num)
            return f'{num:,}{unit} {word}'
    if not value % 1: value = int(value)
    return f'{value:,}{word}'

class Chart:
    def __init__(self, df: pd.DataFrame, date='기준일', close='종가', Open='시가', high='고가', low='저가', volume='거래량', ma=[5, 10, 20, 60, 120, 240],) -> None:
        df.set_index(date)
        for i in ma: df[f'{i}일선'] = df[close].rolling(i).mean()
        self.df, self.date, self.close = (df, date, close)
        self.open, self.high, self.low, self.volume = (Open, high, low, volume)
        self.ma = sorted(ma, reverse=True)
        self.set_fig()
        self.draw()
        return
    def set_fig(self):
        df, close = (self.df, self.close)
        Open, high, low, volume = (self.open, self.high, self.low, self.volume)
        list_ma = self.ma
        fig, ax = plt.subplots(
            4,
            figsize=(12, 6),
            height_ratios=(3, 1, 15, 5)
        )
        fig.subplots_adjust(left=0.01, bottom=0.02, right=0.95, top=0.95, wspace=0, hspace=0)
        ax: list[Axes]
        ax1, ax_none, ax2, ax3 = ax
        ax_none.axis('off')
        m = int(df[high].values.max() * 1.3)
        x = [i for i in range(len(df.index))]
        y = [i for i in range(m)]
        end_y_index = len(y) - 1
        ax1.grid(True, color='#d0d0d0', linewidth=1, zorder=0,)
        for i in list_ma: ax1.plot(x, df[f'{i}일선'], zorder=2,)
        ax1.plot(x, df[self.close].values, color='k', linewidth=2.5, zorder=2,)
        ax1.set_title('Press left mouse button and drag '
                    'to select a region in the top graph')
        xmin, xmax = ax1.get_xlim()
        ax1.set_xlim(xmin, xmax)
        self.asl = ax1.axvspan(xmin, xmax, color=(0, 0, 0, 0.5), zorder=5,)
        self.asr = ax1.axvspan(xmin, xmax, color=(0, 0, 0, 0.5), zorder=5,)
        self.asr.set_visible(False)
        self.axvspan = ax1.axvspan(xmin, xmax, facecolor=(1, 1, 1, 0), zorder=5,)
        self.xmin, self.xmax = (xmin, xmax)
        ax1.set_ylim(df[low].min() * 0.1, df[high].max() * 1.2)
        for a in [ax2, ax3]:
            a.set_facecolor('#fafafa')
            a.grid(True, color='#d0d0d0', linewidth=1, zorder=0,)
        for i in list_ma: ax2.plot(x, df[f'{i}일선'])
        t = time()
        red, red_empty = ([], [])
        blue, blue_empty = ([], [])
        black = []
        pre = None
        for n, (index, i) in enumerate(df.iterrows()):
            c, o, h, l = (i[close], i[Open], i[high], i[low])
            if pre is None: pre = h - l
            if o < c:
                if c < pre: red_empty.append(index)
                else: red.append(index)
            elif c < o:
                if pre < c: blue_empty.append(index)
                else: blue.append(index)
            else: black.append(index)
            pre = c
        t1 = time()
        width_wick = 0.06
        width_line = 0.5
        width_edge = 1.5
        for list_index, body_low, body_high, color_body, color_edge in [
            (red, Open, close, '#fe3032', '#fe3032'),
            (red_empty, Open, close, 'w', '#fe3032'),
            (blue, close, Open, '#0095ff', '#0095ff'),
            (blue_empty, close, Open, 'w', '#0095ff'),
        ]:
            ax2.bar(
                list_index,
                [df[high][i]-df[low][i] for i in list_index],
                width=width_wick,
                bottom=[df[low][i] for i in list_index],
                color=color_edge,
                zorder=2,
            )
            ax2.bar(
                list_index,
                [df[body_high][i]-df[body_low][i] for i in list_index],
                width=width_line,
                bottom=[df[body_low][i] for i in list_index],
                color=color_body,
                edgecolor=color_edge,
                linewidth=width_edge,
                zorder=3, 
            )
        t2 = time()
        ax2.bar(black, [df[high][i]-df[low][i] for i in black], width=width_wick, bottom=[df[low][i] for i in black], color='k', zorder=2,)
        ax2.bar(black, [1 for _ in black], width=width_line, bottom=[df[Open][i] for i in black], color='k', edgecolor='k', linewidth=width_edge, zorder=3,)
        t3 = time()
        width_line = 0.7
        range_volume = df[volume].values.max() * 2
        y_volume = [i for i in range(range_volume)]
        t4 = time()
        ax3.bar(x, df[volume].values, width=width_line, color='#1f77b4', edgecolor='k', linewidth=1, zorder=3)
        t5 = time()
        ax3.set_ylim(0, range_volume)
        for a in [ax1, ax2, ax3]: a.tick_params(left=False, right=True, labelleft=False, labelright=True)
        ax2.set_xticklabels([])
        ax3.set_xticklabels([])
        ax1.yaxis.set_major_formatter(lambda x, _: convert_unit(x))
        ax2.yaxis.set_major_formatter(lambda x, _: convert_unit(x, digit=2))
        ax3.yaxis.set_major_formatter(lambda x, _: convert_unit(x, word='주'))
        t6 = time()
        self.x, self.y, self.end_y_index = (x, y, end_y_index)
        self.y_volume = y_volume
        self.fig, self.ax1, self.ax2, self.ax3 = (fig, ax1, ax2, ax3)
        return
    axvspan = None
    xrange = []
    xsub = 0
    vertical_line = None
    def onselect(self, xmin, xmax, is_move=False):
        if is_move:
            if self.vertical_line is None: self.vertical_line = self.ax1.axvline(color='k', lw=0.8, ls='-', zorder=6)
            self.vertical_line.set_xdata([xmin])
            self.ax1.draw_artist(self.vertical_line)
            return
        t = time()
        if is_move:
            if not self.xrange: return
            elif self.xrange[0] == int(xmin) or self.xrange[1] == int(xmax): pass
            elif self.xsub <= (xmax-xmin): pass
            else: return
        elif (xmax - xmin) < 10: return
        fig, ax1, ax2, ax3 = (self.fig, self.ax1, self.ax2, self.ax3)
        indmin, indmax = (int(xmin), int(xmax))
        self.xrange = [indmin, indmax]
        self.xsub = indmax - indmin - 1
        if not is_move:
            self.asl.set_width(indmin-self.xmin)
            self.asr.set_x(indmax)
            if not self.asr.get_visible(): self.asr.set_visible(True)
            self.axvspan.set_x(xmin)
            self.axvspan.set_width(self.xsub)
        ax2.set_xlim(indmin-1, indmax+2)
        y_min = self.df[self.low].values[max(indmin-1, 0):indmax+1].min() * 0.95
        y_max = self.df[self.high].values[max(indmin-1, 0):indmax+1].max() * 1.03
        ax2.set_ylim(y_min, y_max)
        ax3.set_xlim(indmin-1, indmax+2)
        ax3.set_ylim(0, df[self.volume][max(indmin-1, 0):indmax+1].max() * 1.2)
        fig.canvas.draw()
        fig.canvas.flush_events()
        return
    def aaa(self, xmin, xmax):
        indxmin, indxmax = (int(xmin), int(xmax))
        r = self.x[max(indxmin, 0):min(indxmax, len(self.x))]
        self.ax2.set_autoscalex_on(False)
        self.ax2.set_xbound(r[0], r[1])
        r2 = self.df[self.high].values[max(indxmin, 0):min(indxmax, len(self.x))]
        self.ax2.set_autoscaley_on(False)
        self.ax2.set_ybound(r2[0], r2[1])
        self.fig.canvas.draw()
        return
    def draw(self):
        ax1 = self.ax1
        fc = (0, 0, 0, 0)
        prop = dict(alpha=None, facecolor=fc, edgecolor='#1e78ff', linewidth=4)
        self.span = SpanSelector(
            ax1,
            onselect=self.onselect,
            direction='horizontal',
            useblit=True,
            props=prop,
            interactive=True,
            drag_from_anywhere=True,
            onmove_callback=lambda x1, x2: self.onselect(x1, x2, True),
            grab_range=10,
        )
        self.span._draw_shape(0, len(self.x)-1)
        plt.show()
        return

if __name__ == '__main__':
    list_price = {주가 데이터}
    df = pd.DataFrame(list_price)
    a = Chart(df)