Basic BVH cache. Further implementation requires creation of BVH class.

2 files changed, 157 insertions, 6 deletions

diff --git a/chroma/cache.py b/chroma/cache.py
index fe7d180..392e434 100644
--- a/chroma/cache.py
+++ b/chroma/cache.py
@@ -18,6 +18,12 @@ class GeometryNotFoundError(Exception):
     def __init__(self, msg):
         Exception.__init__(self, msg)
 
+class BVHNotFoundError(Exception):
+    '''A requested bounding volume hierarchy was not found in the
+    on-disk cache.'''
+    def __init__(self, msg):
+        Exception.__init__(self, msg)
+
 def verify_or_create_dir(dirname, exception_msg, logger_msg=None):
     '''Checks if ``dirname`` exists and is a directory.  If it does not exist,
     then it is created.  If it does exist, but is not a directory, an IOError
@@ -39,6 +45,19 @@ class Cache(object):
 
     Use this class to read and write cached geometries or bounding
     volume hierarchies rather than reading and writing disk directly.
+
+    Cached geometries are accessed by a name string.  The name of a
+    geometry should be a legal Python identifier that does not start
+    with an underscore.  (Note that the name string is directly
+    mapped to a filename, so beware of taking Cache names from untrusted
+    sources.  Don't let Little Bobby Tables ruin your day.)
+
+    Bounding volume hierarchies are accessed by MD5 hash of the
+    flattened geometry mesh using the hexadecimal string provided by
+    Mesh.md5().  Multiple BVHs are possible for a given mesh, so the
+    BVH can also be given an optional name (legal Python identifier,
+    no underscore as with geometries).  The version of the BVH
+    with the special name "default" will be the default BVH.
     '''
     
     def __init__(self, cache_dir=os.path.expanduser('~/.chroma/')):
@@ -148,13 +167,70 @@ class Cache(object):
         os.symlink(geo_file, default_geo_file)
 
 
+    def get_bvh_directory(self, mesh_hash):
+        '''Return the full path to the directory corresponding to 
+        ``mesh_hash``.
+        '''
+        return os.path.join(self.bvh_dir, mesh_hash)
+
+    def get_bvh_filename(self, mesh_hash, name='default'):
+        '''Return the full pathname for the BVH file corresponding to 
+        ``name``.
+        '''
+        return os.path.join(self.get_bvh_directory(mesh_hash), name)
+
+
     def list_bvh(self, mesh_hash):
-        pass
+        '''Returns a list the names of all BVHs corresponding to 
+        ``mesh_hash``.
+        '''
+        bvh_dir = self.get_bvh_directory(mesh_hash)
+        if not os.path.isdir(bvh_dir):
+            return []
+        else:
+            return os.listdir(bvh_dir)
+
+    def exist_bvh(self, mesh_hash, name='default'):
+        '''Returns True if a cached BVH exists corresponding to
+        ``mesh_hash`` with the given ``name``.
+        '''
+        return os.path.isfile(self.get_bvh_filename(mesh_hash, name))
+
+    def save_bvh(self, bvh, mesh_hash, name='default'):
+        '''Saves the given chroma.bvh.BVH object to the cache, tagged
+        by the given ``mesh_hash`` and ``name``.
+        '''
+        bvh_dir = self.get_bvh_directory(mesh_hash)
+        verify_or_create_dir(bvh_dir, 
+                             exception_msg='Non-directory already exists '
+                             'where BVH directory should go: ' + bvh_dir)
+        bvh_file = self.get_bvh_filename(mesh_hash, name)
+
+        with open(bvh_file, 'wb') as output_file:
+            pickle.dump(bvh, output_file, 
+                        pickle.HIGHEST_PROTOCOL)
+
+    def load_bvh(self, mesh_hash, name='default'):
+        '''Returns the chroma.bvh.BVH object corresponding to ``mesh_hash``
+        and the given ``name``.  
 
-    def save_bvh(self, mesh_hash, name=None):
-        pass
+        If no BVH exists, raises ``BVHNotFoundError``.
+        '''
+        bvh_file = self.get_bvh_filename(mesh_hash, name)
 
-    def load_bvh(self, mesh_hash, name=None):
-        pass
+        if not os.path.exists(bvh_file):
+            raise BVHNotFoundError(mesh_hash + ':' + name)
 
+        with open(bvh_file, 'rb') as input_file:
+            bvh = pickle.load(input_file)
+            return bvh
 
+    def remove_bvh(self, mesh_hash, name='default'):
+        '''Remove the BVH file associated with ``mesh_hash`` and
+        ``name`` from the cache.
+        
+        If the BVH does not exist, no action is taken.
+        '''
+        bvh_file = self.get_bvh_filename(mesh_hash, name)
+        if os.path.exists(bvh_file):
+            os.remove(bvh_file)
diff --git a/test/test_cache.py b/test/test_cache.py
index e0162c8..e4bb165 100644
--- a/test/test_cache.py
+++ b/test/test_cache.py
@@ -4,7 +4,8 @@ import shutil
 import tempfile
 import binascii
 
-from chroma.cache import verify_or_create_dir, Cache, GeometryNotFoundError
+from chroma.cache import verify_or_create_dir, Cache, GeometryNotFoundError,\
+    BVHNotFoundError
 from chroma.geometry import Geometry, Solid
 from chroma.make import box
 
@@ -167,3 +168,77 @@ class TestCacheGeometry(unittest.TestCase):
 
     def tearDown(self):
         remove_path(self.cache_dir)
+
+class TestCacheBVH(unittest.TestCase):
+    def setUp(self):
+        self.cache_dir = random_tempdir('chroma_cache_test')
+        self.cache = Cache(self.cache_dir)
+
+        self.a = Geometry()
+        self.a.add_solid(Solid(box(1,1,1)))
+        self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10))
+        self.a.flatten()
+
+        self.b = Geometry()
+        self.b.add_solid(Solid(box(2,2,2)))
+        self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10))
+        self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10))
+        self.b.flatten()
+
+        # c is not in cache
+        self.c = Geometry()
+        self.c.add_solid(Solid(box(2,2,2)))
+        self.c.flatten()
+
+        self.a_hash = self.a.mesh.md5()
+        self.b_hash = self.b.mesh.md5()
+        self.c_hash = self.c.mesh.md5()
+
+        self.cache.save_geometry('a', self.a)
+        self.cache.save_geometry('b', self.b)
+
+    def test_list_bvh(self):
+        self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 0)
+        self.cache.save_bvh([], self.a_hash)
+        self.assertIn('default', self.cache.list_bvh(self.a_hash))
+        self.cache.save_bvh([], self.a_hash, 'foo')
+        self.assertIn('foo', self.cache.list_bvh(self.a_hash))
+        self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 2)
+
+    def test_exist_bvh(self):
+        self.cache.save_bvh([], self.a_hash)
+        assert self.cache.exist_bvh(self.a_hash)
+        self.cache.save_bvh([], self.a_hash, 'foo')
+        assert self.cache.exist_bvh(self.a_hash, 'foo')
+        
+    def test_load_bvh_not_found(self):
+        with self.assertRaises(BVHNotFoundError):
+            self.cache.load_bvh(self.c_hash)
+
+        with self.assertRaises(BVHNotFoundError):
+            self.cache.load_bvh(self.a_hash, 'foo')
+
+    def test_save_load_new_bvh(self):
+        self.cache.save_bvh([], self.a_hash)
+        self.cache.load_bvh(self.a_hash)
+        self.cache.save_bvh([], self.a_hash, 'foo')
+        self.cache.load_bvh(self.a_hash, 'foo')
+
+    def test_remove_bvh(self):
+        self.cache.remove_bvh(self.a_hash, 'does_not_exist')
+
+        self.cache.save_bvh([], self.a_hash)
+        self.cache.save_bvh([], self.a_hash, 'foo')
+        assert self.cache.exist_bvh(self.a_hash)
+        assert self.cache.exist_bvh(self.a_hash, 'foo')
+
+        self.cache.remove_bvh(self.a_hash)
+        assert not self.cache.exist_bvh(self.a_hash)
+        assert self.cache.exist_bvh(self.a_hash, 'foo')
+
+        self.cache.remove_bvh(self.a_hash, 'foo')
+        assert not self.cache.exist_bvh(self.a_hash)
+        assert not self.cache.exist_bvh(self.a_hash, 'foo')
+        
+    def tearDown(self):
+        remove_path(self.cache_dir)