Clean up the cache management script.

bin/gpuarray-cache

@@ -1,11 +1,11 @@
 #!/usr/bin/env python
 
 import os
+import sys
 
-def clean(max_size):
+def clean(max_size, path):
     content = []
-    for root, dirs, files in os.walk(os.environ.get('GPUARRAY_CACHE',
-                                                    '~/.gpuarray/cache/')):
+    for root, dirs, files in os.walk(path):
         for file in files:
             fpath = os.path.join(root, file)
             st = os.stat(fpath)
@@ -25,18 +25,18 @@ SUFFIXES = {'B': 1, 'K': 1 << 10, 'M': 1 << 20, 'G': 1 << 30, 'T': 1 << 40,
 
 def get_size(s):
     i = 0
-    while i < len(s) and (s[i].isdigit() or s[i] == '.'):
-        i += 1
-    num = s[:i]
-    suf = s[i:]
+    s = s.strip()
+    if s[-1].upper() in SUFFIXES:
+        num = s[:-1]
+        suf = s[-1].upper()
+    else:
+        num = s
+        suf = ""
     num = float(num)
     if suf != "":
-        letter = suf.strip().upper()
-        if letter not in SUFFIXES:
-            raise ValueError("can't interpret %r" % init)
-        mult = SUFFIXES[letter]
+        mult = SUFFIXES[suf]
     else:
-        mult = 0
+        mult = 1
     return int(num * mult)
 
 
@@ -46,6 +46,10 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='libgpuarray cache maintenance utility')
     parser.add_argument('-s', '--max_size', help='Set the maximum size for pruning (in bytes with suffixes: K, M, G, ...)')
     args = parser.parse_args()
+    path = os.environ.get('GPUARRAY_CACHE_PATH', None)
+    if path is None:
+        print("You need to set GPUARRAY_CACHE_PATH so that this programs knows which path to clean.")
+        sys.exit(1)
 
-    clean(get_size(args.max_size))
+    clean(get_size(args.max_size), path)