58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
|
import sqlite3
|
||
|
|
import argparse
|
||
|
|
from typing import Literal
|
||
|
|
|
||
|
|
KOSPI = "KOSPI"
|
||
|
|
KOSDAQ = "KOSDAQ"
|
||
|
|
|
||
|
|
def queryAllCorpGroup(db: sqlite3.Connection, table: Literal["KOSPI", "KOSDAQ"]) -> list[tuple[str, str]]:
|
||
|
|
assert table in [KOSPI, KOSDAQ]
|
||
|
|
cursor = db.execute("SELECT Code, Name from "+ table)
|
||
|
|
return [ (record[0],record[1]) for record in cursor]
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
|
||
|
|
parser.add_argument("--create", action="store_true", help="create db schema")
|
||
|
|
parser.add_argument("--read", action="store", help="read csv and commit in db")
|
||
|
|
parser.add_argument("--clean", action="store_true", help="clean db")
|
||
|
|
parser.add_argument("--table-name", required=True, choices=["KOSPI", "KOSDAQ"])
|
||
|
|
args = parser.parse_args()
|
||
|
|
if args.create:
|
||
|
|
db = sqlite3.connect("./stock.db")
|
||
|
|
db.execute(f"""
|
||
|
|
CREATE TABLE IF NOT EXISTS "{args.table_name}" (
|
||
|
|
"Code" TEXT,
|
||
|
|
"Name" TEXT,
|
||
|
|
PRIMARY KEY("Code"),
|
||
|
|
FOREIGN KEY("Code") REFERENCES "KRXCorp"("Code")
|
||
|
|
)
|
||
|
|
""")
|
||
|
|
db.close()
|
||
|
|
elif args.read:
|
||
|
|
path = args.read
|
||
|
|
import csv
|
||
|
|
with open(path, "r", encoding='UTF-8') as fp:
|
||
|
|
reader = csv.reader(fp)
|
||
|
|
#next(reader) # skip header
|
||
|
|
# header not exist
|
||
|
|
# collect data
|
||
|
|
data = [ (row[0], row[1]) for row in reader]
|
||
|
|
|
||
|
|
codeDict = { code:name for code, name in data}
|
||
|
|
db = sqlite3.connect("./stock.db")
|
||
|
|
dbData = {code:name for code, name in queryAllCorpGroup(db, args.table_name)}
|
||
|
|
# remove common code
|
||
|
|
commonCode = set(dbData.keys()) & set(codeDict.keys())
|
||
|
|
for code in commonCode:
|
||
|
|
del codeDict[code]
|
||
|
|
del dbData[code]
|
||
|
|
db.execute("BEGIN")
|
||
|
|
# insert new code
|
||
|
|
db.executemany(f"INSERT INTO {args.table_name} (Code, Name) VALUES (?, ?)", codeDict.items())
|
||
|
|
# delete old code
|
||
|
|
db.executemany(f"DELETE FROM {args.table_name} WHERE code = ?", dbData.keys())
|
||
|
|
db.execute("COMMIT")
|
||
|
|
db.commit()
|
||
|
|
db.close()
|