Add update_codes and script to cli (#49)

mlouielu · Oct 14, 2018 · 43cc049 · 43cc049
1 parent 9542a13
commit 43cc049
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git a/twstock/__init__.py b/twstock/__init__.py
@@ -7,8 +7,8 @@
 from twstock import realtime
 
 from twstock.analytics import BestFourPoint
-from twstock.codes import twse, tpex, codes
+from twstock.codes import __update_codes, twse, tpex, codes
 from twstock.stock import Stock
 
 
-__version__ = '1.1.1'
+__version__ = '1.1.2'
diff --git a/twstock/cli/__init__.py b/twstock/cli/__init__.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import argparse
+from twstock.codes import __update_codes
 from twstock.cli import best_four_point
 from twstock.cli import stock
 from twstock.cli import realtime
@@ -12,6 +13,7 @@ def run():
     parser.add_argument('-b', '--bfp', nargs='+')
     parser.add_argument('-s', '--stock', nargs='+')
     parser.add_argument('-r', '--realtime', nargs='+')
+    parser.add_argument('-U', '--upgrade-codes', action='store_true', help='Update entites codes')
     args = parser.parse_args()
 
     if args.bfp:
@@ -20,3 +22,7 @@ def run():
         stock.run(args.stock)
     elif args.realtime:
         realtime.run(args.realtime)
+    elif args.upgrade_codes:
+        print('Start to update codes')
+        __update_codes()
+        print('Done!')
diff --git a/twstock/codes/__init__.py b/twstock/codes/__init__.py
@@ -0,0 +1,2 @@
+from .fetch import __update_codes
+from .codes import codes, tpex, twse
diff --git a/twstock/codes.py → twstock/codes/codes.py b/twstock/codes.py → twstock/codes/codes.py
@@ -14,8 +14,8 @@
 ROW = namedtuple('StockCodeInfo', ['type', 'code', 'name', 'ISIN', 'start',
                                    'market', 'group', 'CFI'])
 PACKAGE_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
-TPEX_EQUITIES_CSV_PATH = os.path.join(PACKAGE_DIRECTORY, 'codes', 'tpex_equities.csv')
-TWSE_EQUITIES_CSV_PATH = os.path.join(PACKAGE_DIRECTORY, 'codes', 'twse_equities.csv')
+TPEX_EQUITIES_CSV_PATH = os.path.join(PACKAGE_DIRECTORY, 'tpex_equities.csv')
+TWSE_EQUITIES_CSV_PATH = os.path.join(PACKAGE_DIRECTORY, 'twse_equities.csv')
 
 codes = {}
 tpex = {}

diff --git a/twstock/codes/fetch.py b/twstock/codes/fetch.py
@@ -6,6 +6,7 @@
 # TPEx equities = 上櫃證券
 #
 
+import os
 import csv
 from collections import namedtuple
 
@@ -52,6 +53,13 @@ def to_csv(url, path):
             writer.writerow([_ for _ in d])
 
 
+def __update_codes():
+    def get_directory():
+        return os.path.dirname(os.path.abspath(__file__))
+    to_csv(TWSE_EQUITIES_URL, os.path.join(get_directory(), 'twse_equities.csv'))
+    to_csv(TPEX_EQUITIES_URL, os.path.join(get_directory(), 'tpex_equities.csv'))
+
+
 if __name__ == '__main__':
     to_csv(TWSE_EQUITIES_URL, 'twse_equities.csv')
     to_csv(TPEX_EQUITIES_URL, 'tpex_equities.csv')