Stock/meta.py

58 lines
2.1 KiB
Python
Raw Normal View History

2023-07-22 16:10:40 +09:00
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()