make-quant-py/streamlit_quant/quantcommon.py

import os
from urllib.parse import quote_plus

import pandas as pd
import pymysql
from dotenv import load_dotenv
from sqlalchemy import create_engine


class QuantCommon:
    def __init__(self):
        load_dotenv()
        self.user = os.getenv('DB_USER')
        self.pw = os.getenv('DB_PW')
        self.engine_for_pw = quote_plus(self.pw)
        self.host = os.getenv('DB_HOST')
        self.port = int(os.getenv('DB_PORT'))
        self.db = os.getenv('DB_DB')

    def create_engine(self):
        return create_engine(f'mysql+pymysql://{self.user}:{self.engine_for_pw}@{self.host}:{self.port}/{self.db}')

    def connect(self):
        return pymysql.connect(user=self.user,
                      passwd=self.pw,
                      host=self.host,
                      port=self.port,
                      db=self.db,
                      charset='utf8')

    def get_ticker_list(self):
        engine = self.create_engine()

        try:
            ticker_list = pd.read_sql("""
                        select * from kor_ticker
                        where 기준일 = (select max(기준일) from kor_ticker) 
                            and 종목구분 = '보통주';
                """, con=engine)
        finally:
            engine.dispose()

        return ticker_list

    def get_value_list(self):
        engine = self.create_engine()

        try:
            value_list = pd.read_sql("""
            select * from kor_value
            where 기준일 = (select max(기준일) from kor_value);
            """, con=engine)
        finally:
            engine.dispose()

        return value_list

    def get_price_list(self, interval_month):
        engine = self.create_engine()

        try:
            price_list = pd.read_sql(f"""
            select 날짜, 종가, 종목코드
            from kor_price
            where 날짜 >= (select (select max(날짜) from kor_price) - interval {interval_month} month);
            """, con=engine)
        finally:
            engine.dispose()

        return price_list

    def get_price_list_by_code(self, codes):
        engine = self.create_engine()

        try:
            price_list = pd.read_sql(f"""
            select * from kor_price
            where 종목코드 in ({codes});
            """, con=engine)
        finally:
            engine.dispose()

        return price_list

    def get_fs_list(self):
        engine = self.create_engine()

        try:
            fs_list = pd.read_sql("""
            select * from kor_fs
            where 계정 in ('당기순이익', '매출총이익', '영업활동으로인한현금흐름', '자산', '자본')
            and 공시구분 = 'q';
            """, con=engine)
        finally:
            engine.dispose()

        return fs_list

    def get_fs_list_by_account_and_date(self, account, date):
        engine = self.create_engine()

        try:
            fs_list = pd.read_sql(f"""
            select * from kor_fs
            where 계정 in ({account})
            and 기준일 in ({date})
            and 공시구분 = 'y';
            """, con=engine)
        finally:
            engine.dispose()

        return fs_list

    def get_expanded_fs_list(self):
        engine = self.create_engine()

        try:
            fs_list = pd.read_sql("""
            select * from kor_fs
            where 계정 in ('매출액', '당기순이익', '법인세비용', '이자비용', '현금및현금성자산',
                '부채', '유동부채', '유동자산', '비유동자산', '감가상각비')
            and 공시구분 = 'q';
            """, con=engine)
        finally:
            engine.dispose()

        return fs_list

    def get_sector_list(self):
        engine = self.create_engine()

        try:
            sector_list = pd.read_sql("""
            select * from kor_sector
            where 기준일 = (select max(기준일) from kor_sector);	
            """, con=engine)
        finally:
            engine.dispose()

        return sector_list