Stock/gen.py
2022-06-19 17:58:53 +09:00

215 lines
8.1 KiB
Python

import argparse
import json
import os
import sqlite3
from typing import Dict, List
from render import *
import db as database
from jinja2 import Environment, PackageLoader, select_autoescape
import pandas as pd
import tqdm
class DataStore:
def __init__(self) -> None:
self.db = sqlite3.connect("stock.db")
self.pricesCache: Dict[str,] = {}
def getAllKRXCorp(self) -> List[database.KRXCorp]:
return database.GetAllKRXCorp(self.db)
def getStockPrice(self,code,length) -> pd.DataFrame:
if code in self.pricesCache and len(self.pricesCache[code]) >= length:
return self.pricesCache[code]
else:
s = GetStockPriceFrom(self.db,code,length)
s = pd.DataFrame(s,
columns=[s for s in database.STOCK_INDEX.__members__.keys()])
s.set_index("DATE", inplace=True)
self.pricesCache[code] = s
return self.pricesCache[code]
def clearCache(self) -> None:
self.pricesCache = {}
def __del__(self) -> None:
self.db.close()
class OutputCollectorElement:
def __init__(self, name: str, description: str) -> None:
self.name = name
self.description = description
self.corpListByDate:Dict[str,database.KRXCorp] = {}
def __str__(self) -> str:
return f"OutputCollectorElement:{self.name}"
def addCorp(self, date, corp):
self.corpListByDate.setdefault(date, []).append(corp)
def toDict(self) -> Dict:
return {
"name": self.name,
"description": self.description,
"corpListByDate": {k:[d.toDict() for d in v]
for k,v in self.corpListByDate.items()}
}
class OutputCollector:
def __init__(self) -> None:
self.data: Dict[str,OutputCollectorElement] = {}
def addResult(self, key, help = ""):
"""
add output category to collect
"""
self.data[key] = OutputCollectorElement(key, help)
def collect(self, key, corp, date):
self.data[key].addCorp(date, corp)
def isVolumeNTimes(stock: pd.DataFrame, mul: float, nday:int, order=1) -> bool:
return stock.iloc[nday]['VOLUME'] > stock.iloc[nday+order]['VOLUME'] * mul
def isVolumeMulPriceGreaterThan(stock: pd.DataFrame, threshold: int, nday: int) -> bool:
return stock.iloc[nday]['VOLUME'] * stock.iloc[nday]['CLOSE'] > threshold
def isMACDCrossSignal(signal: pd.Series, macd: pd.Series, nday: int, order=1) -> bool:
return (signal.iloc[nday] > macd.iloc[nday] and
signal.iloc[nday+order] < macd.iloc[nday+order])
def isRelativeDiffLessThan(a:pd.Series,b:pd.Series, threshold: float,nday:int) -> bool:
return abs(a.iloc[nday] - b.iloc[nday]) / b.iloc[nday] < threshold
def isDiffGreaterThan(a:pd.Series,b:pd.Series, nday:int) -> bool:
"""a is bigger than b"""
return (a.iloc[nday] > b.iloc[nday])
def prepareCollector(collector: OutputCollector) -> None:
import pages
for item in pages.GenLists:
collector.addResult(item["name"], item["description"])
def every(f, xs):
for x in xs:
if not f(x):
return False
return True
def collect(data: DataStore, collector: OutputCollector, corp: database.KRXCorp
, nday: int) -> None:
stock = data.getStockPrice(corp.Code,120)
if len(stock) < 120:
return
if (stock.iloc[nday]['VOLUME'] <= 0):
return
close = stock["CLOSE"]
d5 = stock["CLOSE"].loc[::-1].rolling(window=5
).mean().dropna().loc[::-1]
d10 = stock["CLOSE"].loc[::-1].rolling(window=10
).mean().dropna().loc[::-1]
d20 = stock["CLOSE"].loc[::-1].rolling(window=20
).mean().dropna().loc[::-1]
d30 = stock["CLOSE"].loc[::-1].rolling(window=30
).mean().dropna().loc[::-1]
d60 = stock["CLOSE"].loc[::-1].rolling(window=60
).mean().dropna().loc[::-1]
a = [d5, d10, d20, d30, d60]
if every(lambda i: isRelativeDiffLessThan(i,close,0.05,nday), a):
collector.collect("뭉침", corp, stock.index[nday])
if every(lambda i: isRelativeDiffLessThan(i,close,0.01,nday), a):
collector.collect("뭉침01", corp, stock.index[nday])
if every(lambda i: isRelativeDiffLessThan(i,close,0.03,nday), a):
collector.collect("뭉침03", corp, stock.index[nday])
if (isRelativeDiffLessThan(d5, d20, 0.01, nday) and
isRelativeDiffLessThan(d5, d60, 0.01, nday)):
collector.collect("cross 2", corp, stock.index[nday])
if (isVolumeNTimes(stock, 3, 0) and
isVolumeMulPriceGreaterThan(stock, 100000, nday)):
collector.collect("cross 3", corp, stock.index[nday])
if (isRelativeDiffLessThan(d20, d60, 0.01, nday) and
isVolumeMulPriceGreaterThan(stock, 1000000, nday)):
collector.collect("cross 4", corp, stock.index[nday])
if (isDiffGreaterThan(d5, d20, nday)):
collector.collect("d20d5", corp, stock.index[nday])
if (isVolumeNTimes(stock, 5, nday)):
collector.collect("d20d5VolumeX5", corp, stock.index[nday])
if (isRelativeDiffLessThan(d5, d20, 0.03, nday) and
isRelativeDiffLessThan(d5, d60, 0.03, nday) and
isVolumeNTimes(stock, 3, nday)):
collector.collect("DiffDistance", corp, stock.index[nday])
if (isVolumeNTimes(stock, 3, nday)):
collector.collect("volume", corp, stock.index[nday])
if (isVolumeMulPriceGreaterThan(stock, 50000000, nday)):
collector.collect("volume5", corp, stock.index[nday])
if (isVolumeNTimes(stock, 5, nday)):
collector.collect("volumeX5", corp, stock.index[nday])
ewm5 = stock["CLOSE"].loc[::-1].ewm(span=5).mean().loc[::-1]
ewm10 = stock["CLOSE"].loc[::-1].ewm(span=10).mean().loc[::-1]
macd = (ewm5 - ewm10)
signal = macd.loc[::-1].ewm(span=4).mean().loc[::-1]
if (isMACDCrossSignal(macd, signal, nday)):
collector.collect("macd", corp, stock.index[nday])
parser = argparse.ArgumentParser(description="주식 검색 정보를 출력합니다.")
parser.add_argument("--format", "-f", choices=["json", "html"], default="html",
help="출력 포맷을 지정합니다. 기본값은 html입니다.")
parser.add_argument("--dir", "-d", default=".", help="출력할 폴더를 지정합니다.")
parser.add_argument("--corp", "-c", help="주식 코드를 지정합니다. 지정하지 않으면 모든 주식을 검색합니다.")
parser.add_argument("--printStdout", action="store_true", help="출력한 결과를 표준 출력으로 출력합니다.")
parser.add_argument("--version", "-v", action="version", version="%(prog)s 1.0")
parser.add_argument("--verbose", "-V", action="store_true", help="출력할 내용을 자세히 표시합니다.")
if __name__ == "__main__":
args = parser.parse_args()
dataStore = DataStore()
krx_corps = dataStore.getAllKRXCorp()
if args.corp:
krx_corps = [corp for corp in krx_corps if corp.Code == args.corp]
env = Environment(
loader=PackageLoader('render', 'templates'),
autoescape=select_autoescape(['html', 'xml'])
)
collector = OutputCollector()
prepareCollector(collector)
for corp in tqdm.tqdm(krx_corps):
for nday in range(0, 5):
collect(dataStore, collector, corp, nday)
dataStore.clearCache()
for k,v in collector.data.items():
if args.format == "json":
data = json.dumps(v.toDict(), indent=4, ensure_ascii=False)
if args.printStdout:
print(k)
print(data)
else:
with open(os.path.join(args.dir, k + ".json", encoding="UTF-8"), "w") as f:
f.write(data)
else:
template = env.get_template("Lists.html")
days = v.corpListByDate.keys()
days = list(days)
days.sort(reverse=True)
days = days[:5]
html = template.render(collected=v, title=k, days=days, lastUpdate=datetime.date.today().isoformat())
if args.printStdout:
print(html)
else:
with open(os.path.join(args.dir, k + ".html"), "w", encoding="UTF-8") as f:
f.write(html)