Added disk cache and -k option to refresh

ctabus.py

@@ -1,5 +1,6 @@
 from urllib.parse import urlencode
 from urllib.request import urlopen
+from disk_cache import disk_cache
 import json
 from sensitive import api
 
@@ -10,7 +11,7 @@ def get_data(type, api_key=api, timeout=None, **args):
     args['format'] = 'json'
     url = base_url.format(type=type, query=urlencode(args))
     if timeout is not None:
-        response = urlopen(url,timeout = timeout)
+        response = urlopen(url, timeout=timeout)
     else:
         response = urlopen(url)
     data = json.load(response)['bustime-response']
@@ -25,13 +26,16 @@ def get_times(stop_id, api_key=api, timeout=None):
     return get_data('getpredictions', api_key, stpid=stop_id, timeout=timeout)
 
 
+@disk_cache
 def get_routes(api_key=api, timeout=None):
     return get_data('getroutes', api_key, timeout=timeout)
 
 
+@disk_cache
 def get_directions(route, api_key=api, timeout=None):
     return get_data('getdirections', api_key, rt=route, timeout=timeout)
 
 
+@disk_cache
 def get_stops(route, direction, api_key=api, timeout=None):
     return get_data('getstops', api_key, rt=route, dir=direction, timeout=timeout)

disk_cache.py

@@ -13,11 +13,15 @@ def make_key(*args, **kwargs):
 
 
 class disk_cache:
+    caches = []
+    """Decorator to make a function with lru_cache that can be written to disk"""
+
     def __init__(self, func):
         self.fname = "{}.{}.dc".format(func.__module__, func.__name__)
         self.fname = os.path.join(cache_path, self.fname)
         self.func = func
         self.load_cache()
+        disk_cache.caches.append(self)
 
     def __call__(self, *args, **kwargs):
         key = make_key(*args, **kwargs)
@@ -40,7 +44,7 @@ class disk_cache:
 
     def save_cache(self):
         with lzma.open(self.fname, 'wb') as file:
-            pickle.dump(self.cache, file)
+            pickle.dump(self.cache, file, pickle.HIGHEST_PROTOCOL)
 
     def delete_cache(self):
         os.remove(self.fname)

main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 from dateutil.parser import parse as date_parse
 from dateutil import tz
+from disk_cache import disk_cache
 import argparse
 import ctabus
 import datetime
@@ -195,7 +196,14 @@ if __name__ == '__main__':
                         type=int, help='checks periodically')
     parser.add_argument('-r', '--route', default=None)
     parser.add_argument('-d', '--direction', default=None)
+    parser.add_argument('-k', '--kill-cache', action="store_true")
     parser.add_argument('arg', nargs='+', metavar='(stop-id | cross streets)')
     args = parser.parse_args()
     sys.stderr = open(osp.join(osp.dirname(__file__), 'stderr.log'), 'w')
+    if args.kill_cache:
+        for cache_obj in disk_cache.caches:
+            cache_obj.kill_cache()
     main(args)
+    for cache_obj in disk_cache.caches:
+        if cache_obj.fresh:
+            cache_obj.save_cache()