Change Cachedir.open to allow dict_key instead of basename, and optional suffix

2 files changed, 25 insertions, 4 deletions

diff --git a/morphlib/cachedir.py b/morphlib/cachedir.py
index d0763c4c..6bd75608 100644
--- a/morphlib/cachedir.py
+++ b/morphlib/cachedir.py
@@ -64,7 +64,7 @@ class CacheDir(object):
 
         return os.path.join(self.dirname, key + suffix)
 
-    def open(self, relative_name):
+    def open(self, relative_name_or_cache_key, suffix=''):
         '''Open a file for writing in the cache.
         
         The file will be written with a temporary name, and renamed to
@@ -75,7 +75,11 @@ class CacheDir(object):
         temporary file.
         
         '''
-        
-        path = os.path.join(self.dirname, relative_name)
+
+        if type(relative_name_or_cache_key) is dict:
+            path = self.name(relative_name_or_cache_key)
+        else:
+            path = os.path.join(self.dirname, relative_name_or_cache_key)
+        path += suffix
         return morphlib.savefile.SaveFile(path, 'w')
 
diff --git a/morphlib/cachedir_tests.py b/morphlib/cachedir_tests.py
index f6cccc97..e7ab9354 100644
--- a/morphlib/cachedir_tests.py
+++ b/morphlib/cachedir_tests.py
@@ -97,12 +97,29 @@ class CacheDirTests(unittest.TestCase):
         pathname = self.cachedir.name(dict_key)
         self.assert_(pathname.startswith(self.cachedir.dirname + '/'))
 
-    def test_allows_file_to_be_written(self):
+    def test_allows_file_to_be_written_via_basename(self):
         f = self.cachedir.open('foo')
         f.write('bar')
         f.close()
         self.assertEqual(self.cat('foo'), 'bar')
 
+    def test_allows_file_to_be_written_via_basename_and_suffix(self):
+        f = self.cachedir.open('foo', '.blip')
+        f.write('bar')
+        f.close()
+        self.assertEqual(self.cat('foo.blip'), 'bar')
+
+    def test_allows_file_to_be_written_via_dict_key(self):
+        dict_key = {
+            'kind': 'chunk',
+            'meh': 'moo',
+        }
+        name = self.cachedir.name(dict_key)
+        f = self.cachedir.open(dict_key)
+        f.write('bar')
+        f.close()
+        self.assertEqual(self.cat(name), 'bar')
+
     def test_allows_file_to_be_aborted(self):
         f = self.cachedir.open('foo')
         f.write('bar')