Updated file caching to use both accession & file_type.

chipathlon/db.py

@@ -6,6 +6,7 @@ import traceback
 import os
 import itertools
 import time
+import collections
 import chipathlon.conf
 from pprint import pprint
 import hashlib
@@ -35,7 +36,7 @@ class MongoDB(object):
         self.password = password
         self.client = MongoClient(host)
         self.db = self.client.chipseq
-        self.cache = {}
+        self.cache = collections.defaultdict(dict)
         try:
             self.db.authenticate(username, password, mechanism="SCRAM-SHA-1")
         except:
@@ -45,12 +46,12 @@ class MongoDB(object):
         self.gfs = gridfs.GridFS(self.db)
         return
 
-    def add_cache(self, function, key, data):
+    def add_cache(self, accession, file_type, data):
         """
-        :param function: The function name
-        :type function: str
-        :param key: The key to add the cache entry under.
-        :type key: Any hashable
+        :param accession: The accession of the file to store.
+        :type accession: str
+        :param file_type: The type of file to store.
+        :type file_type: str
         :param data: The data to add to the cache.
         :type data: Object
 
@@ -60,23 +61,20 @@ class MongoDB(object):
         peak calling tools.  In these cases we don't want to request info
         from the database multiple times for the same data.
         """
-        if function not in self.cache:
-            self.cache[function] = {}
-        self.cache[function][key] = data
+        self.cache[accession][file_type] = data
         return
 
-    def get_cache(self, function, key):
+    def get_cache(self, accession, file_type):
         """
-        :param function: The function name
-        :type function: str
-        :param key: The key to get from the cache.
-        :type key: Any hashable
+        :param accession: The accession of the file to retrieve.
+        :type accession: str
+        :param file_type: The type of file to retrieve.
+        :type file_type: str
 
         Gets a data item from the internal cache.
         """
-        if function in self.cache:
-            if key in self.cache[function]:
-                return self.cache[function][key]
+        if accession in self.cache:
+            return self.cache[accession].get(file_type)
         return None
 
     def delete_result(self, result, genome):
@@ -331,7 +329,7 @@ class MongoDB(object):
         valid = True
         msg = ""
         data = {}
-        check_cache = self.get_cache("get_sample", accession)
+        check_cache = self.get_cache(accession, file_type)
         if check_cache is not None:
             msg = "Retrieved data from cache."
             data = check_cache
@@ -342,7 +340,7 @@ class MongoDB(object):
             })
             if cursor.count() == 1:
                 data = cursor.next()
-                self.add_cache("get_sample", accession, data)
+                self.add_cache(accession, file_type, data)
             else:
                 valid = False
                 msg = "Found %s files with accession: %s, file_type: %s. Should only be 1." % (