2-2. DB에서 시세 조회하는 API 만들기

2021. 6. 6. 19:54Financial Analysis

2-1장에서는 웹 페이지에서 기업들의 일별 시세를 스크래핑한 뒤, MariaDB에 저장하는 기능을 구현하였다.

이번 장에서는 DB에 저장된 종목들의 일별 시세를 기간 별로 조회할 수 있는 API를 만들어보자.

import pymysql
import pandas as pd
from datetime import datetime
from datetime import timedelta
import re

class MarketDB:
    def __init__(self):
        """생성자 : MarketDB 연결 및 종목코드 딕셔너리 생성"""
        
    def __del__(self):
        """소멸자 : MariaDB 연결 해제"""

    def get_comp_info(self):
        """company_info 테이블에서 읽어와서 codes에 저장"""

    def get_daily_price(self, code, start_date=None, end_date=None):
        """KRX 종목별 시세를 데이터프레임 형태로 변환"""

위는 시세 조회를 담당하는 MarketDB 클래스의 구조를 나타내고 있다.

    def __init__(self):
        """생성자 : MarketDB 연결 및 종목코드 딕셔너리 생성"""
        self.conn = pymysql.connect(host='localhost', user='root', password='3214', db='Investar', charset='utf8')
        self.codes = {}
        self.get_comp_info()
        
    def __del__(self):
        """소멸자 : MariaDB 연결 해제"""
        self.conn.close()


생성자에서는 클래스 내부적으로 자체적인 리터럴 딕셔너리를 가지고, get_comp_info 함수를 호출하여 딕셔너리를 채운다.

소멸자가 호출되면 DB와의 연결을 해제한다.

    def get_comp_info(self):
        """company_info 테이블에서 읽어와서 codes에 저장"""
        sql = 'SELECT * FROM company_info'
        krx = pd.read_sql(sql, self.conn)
        for idx in range(len(krx)):
            self.codes[krx['code'].values[idx]] = krx['company'].values[idx]

get_comp_info 함수에서는 DB의 company_info 테이블의 데이터를 가져와서 클래스 내부의 codes 딕셔너리에

복사해준다.

   def get_daily_price(self, code, start_date=None, end_date=None):
        """KRX 종목별 시세를 데이터프레임 형태로 변환"""
        if (start_date is None):
            one_year_age = datetime.today() - timedelta(days=365)
            start_date = one_year_age.strftime('%Y-%m-%d')
            print("start date is initialized to '{}'".format(start_date))
        else:
            start_lst = re.split('\D+', start_date)
            if (start_lst[0] == ''):
                start_lst = int(start_lst[1:])
            start_year = int(start_lst[0])
            start_month = int(start_lst[1])
            start_day = int(start_lst[2])
            if start_year < 1900 or start_year> 2200:
                print(f"ValueError : start_year({start_year:d}) is unprocessable.")
                return
            if start_month < 1 or start_month > 12:
                print(f"ValueError : start_month({start_month:d}) is unprocessable.")
                return
            if start_day < 1 or start_day > 31:
                print(f"ValueError : start_day({start_day:d}) is unprocessable.")
                return
            start_date = f"{start_year:04d}-{start_month:02d}-{start_day:02d}"
                                          .
                                          .
                                          .

get_daily_price 함수가 이 클래스의 핵심 함수인데 길이가 상당히 길기 때문에 하나하나 살펴보자.

먼저 인수로 시세 조회할 '종목의 코드'와 일별 시세의 '시작일'과 '종료일'을 받는다.

만약 시작일을 입력하지 않았다면 지금으로부터 1년 전의 날짜를 시작일로 지정한다.

시작일이 주어졌다면 정규식을 이용하여 parsing한 결과를 년, 월, 일에 저장한다.

시작일이 가능한 변수 범위를 벗어났다면 예외처리하는 코드도 추가한다. 종료일 역시 마찬가지로 구현한다.

        codes_keys = list(self.codes.keys())
        codes_values = list(self.codes.values())
        if code in codes_keys:
            pass
        elif code in codes_values:
            idx = codes_values.index(code)
            code = codes_keys[idx]
        else:
            print("ValueError : Code({}) doesn't exist.".format(code))

        sql = f"SELECT * FROM daily_price WHERE code = '{code}'" \
              f"and date >= '{start_date}' and date <= '{end_date}'"
        df = pd.read_sql(sql, self.conn)
        df.index = df['date']
        return df

시세 조회 API에서는 조회할 종목을 조회할 때 종목코드와 회사명 중 어느 것도 입력할 수 있도록 지원한다.

딕셔너리 특성 상 값(회사명)을 통해 키(종목코드)를 조회할 수 없기 때문에 먼저 codes 딕셔너리에서

Key와 Value를 따로 분리하여 list에 저장한다. 입력받은 인수가 Values(회사명)에 존재한다면

해당값이 list에 몇번째 인덱스에 존재하는지 조회하여 해당 인덱스의 종목 모드를 알아낸다.

알아낼 정보는 모두 알아냈으므로 sql문을 통해 DB에서 조회하여 데이터프레임으로 반환한다.

>>> from Investar.MarketDB import MarketDB
>>> mDB = MarketDB()
>>> mDB.get_daily_price("CJ대한통운", '2020.01.01', '2020-12-12')
              code        date    open    high     low   close  diff  volume
date
2020-01-02  000120  2020-01-02  155000  155500  152500  152500  2500   24403
2020-01-03  000120  2020-01-03  153500  153500  147000  148000  4500   72998
2020-01-06  000120  2020-01-06  147500  149000  144500  147000  1000   67028
2020-01-07  000120  2020-01-07  147000  147500  145000  147000     0   25991
2020-01-08  000120  2020-01-08  144500  145500  141000  143500  3500   46281
...            ...         ...     ...     ...     ...     ...   ...     ...
2020-12-07  000120  2020-12-07  159500  160000  157500  158000  3500  112933
2020-12-08  000120  2020-12-08  159000  160500  158000  159500  1500   70202
2020-12-09  000120  2020-12-09  159000  161000  158500  160500  1000   76452
2020-12-10  000120  2020-12-10  162000  162000  158500  161500  1000  135637
2020-12-11  000120  2020-12-11  161500  162500  161000  161000   500   75168

조회된 정보를 바탕으로 종가 그래프를 그려보자.

>>> fig = plt.figure()
>>> ax1 = fig.add_subplot(2,1,1)
>>> ax1.plot(cj_dil['date'], cj_dil['close'])
[<matplotlib.lines.Line2D object at 0x000001B5636291F0>]
>>> ax1.grid(True)
>>> ax2 = fig.add_subplot(2,1,2)
>>> ax2.bar(cj_dil['date'], cj_dil['volume'])
<BarContainer object of 236 artists>
>>> fig.suptitle("CJ logistics Stock")
Text(0.5, 0.98, 'CJ logistics Stock')
>>> ax2.set_title("CJ logistics Volume")
Text(0.5, 1.0, 'CJ logistics Volume')
>>> fig.tight_layout()
>>> fig.show()
CJ대한통운 주가 그래프