From 5437b03cf90c61514e919d902f24182532f0ff8d Mon Sep 17 00:00:00 2001
From: Raphael Roberts <raphael.roberts48@gmail.com>
Date: Mon, 1 Apr 2019 13:57:59 -0500
Subject: [PATCH] Added disk cache and -k option to refresh

---
 ctabus.py     | 6 +++++-
 disk_cache.py | 6 +++++-
 main.py       | 8 ++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/ctabus.py b/ctabus.py
index 4ef42ea..8491553 100644
--- a/ctabus.py
+++ b/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)
diff --git a/disk_cache.py b/disk_cache.py
index 8e97d54..a5e84d4 100644
--- a/disk_cache.py
+++ b/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)
diff --git a/main.py b/main.py
index e5e41fe..d867493 100755
--- a/main.py
+++ b/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()