From 19616b59cb124c62c886e5e3bf47de3a49905b48 Mon Sep 17 00:00:00 2001
From: rainborntw <rainborntw@gmail.com>
Date: Sat, 30 Oct 2021 14:09:28 +0800
Subject: [PATCH] Add cache to Stock

---
 twstock/stock.py | 39 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/twstock/stock.py b/twstock/stock.py
index 7c926c2..ca6a4ad 100644
--- a/twstock/stock.py
+++ b/twstock/stock.py
@@ -23,6 +23,7 @@
     import analytics
     from codes import codes
 
+import os
 
 TWSE_BASE_URL = 'http://www.twse.com.tw/'
 TPEX_BASE_URL = 'http://www.tpex.org.tw/'
@@ -142,16 +143,19 @@ def purify(self, original_data):
 
 class Stock(analytics.Analytics):
 
-    def __init__(self, sid: str, initial_fetch: bool=True):
+    cache_dir = os.path.expanduser('~/.twstock/')
+    def __init__(self, sid: str):
         self.sid = sid
         self.fetcher = TWSEFetcher(
         ) if codes[sid].market == '上市' else TPEXFetcher()
         self.raw_data = []
         self.data = []
+        self.cache_filename = self.cache_dir + self.sid
 
         # Init data
-        if initial_fetch:
+        if not self.read_from_cache():
             self.fetch_31()
+            self.update_cache()
 
     def _month_year_iter(self, start_month, start_year, end_month, end_year):
         ym_start = 12 * start_year + start_month - 1
@@ -184,6 +188,37 @@ def fetch_31(self):
         self.data = self.data[-31:]
         return self.data
 
+    def read_from_cache(self):
+        if os.path.isfile(self.cache_filename):
+            # check if mtime is today
+            today = datetime.datetime.now().date()
+            if datetime.datetime.fromtimestamp(os.path.getmtime(self.cache_filename)).date() != today:
+                return False
+
+            with open(self.cache_filename, 'r') as f:
+                while True:
+                    data = f.readline().strip().split(',')
+                    if data[0] == '':
+                        break
+
+                    data[0] = data[0][:-9].replace('-', '/')
+                    self.data.append(self.fetcher._make_datatuple(data))
+
+            return True
+
+        return False
+
+    def update_cache(self):
+        if not os.path.exists(self.cache_dir):
+            os.mkdir(self.cache_dir)
+
+        with open(self.cache_filename, 'w') as f:
+            for d in self.data:
+                data = []
+                for val in d:
+                    data.append(str(val))
+                f.write(','.join(data) + '\n')
+
     @property
     def date(self):
         return [d.date for d in self.data]