Include only files matching regexps, removes them from after adding them

2 files changed, 60 insertions, 11 deletions

diff --git a/morphlib/bins.py b/morphlib/bins.py
index ed2ac555..993ae3ad 100644
--- a/morphlib/bins.py
+++ b/morphlib/bins.py
@@ -22,18 +22,67 @@ Binaries are chunks, strata, and system images.
 
 
 import logging
+import os
+import re
 import tarfile
 
 import morphlib
 
 
-def create_chunk(rootdir, chunk_filename, globs):
-    '''Create a chunk from the contents of a directory.'''
-    logging.debug('Creating chunk file %s from %s' % (chunk_filename, rootdir))
+def create_chunk(rootdir, chunk_filename, regexps):
+    '''Create a chunk from the contents of a directory.
+    
+    Only files and directories that match at least one of the regular
+    expressions are accepted. The regular expressions are implicitly
+    anchored to the beginning of the string, but not the end. The 
+    filenames are relative to rootdir.
+    
+    '''
+    
+    def mkrel(filename):
+        assert filename.startswith(rootdir)
+        if filename == rootdir:
+            return '.'
+        assert filename.startswith(rootdir + '/')
+        return filename[len(rootdir + '/'):]
+
+    def matches(filename):
+        return any(x.match(filename) for x in compiled)
+
+    def names_to_root(filename):
+        yield filename
+        while filename != rootdir:
+            filename = os.path.dirname(filename)
+            yield filename
+
+    logging.debug('Creating chunk file %s from %s with regexps %s' % 
+                    (chunk_filename, rootdir, regexps))
+
+    compiled = [re.compile(x) for x in regexps]
+    include = set()
+    for dirname, subdirs, basenames in os.walk(rootdir):
+        if matches(dirname):
+            include.add(dirname)
+        filenames = [os.path.join(dirname, x) for x in basenames]
+        for filename in filenames:
+            if matches(mkrel(filename)):
+                for name in names_to_root(filename):
+                    include.add(name)
+
+    include = sorted(include)
+
     tar = tarfile.open(name=chunk_filename, mode='w:gz')
-    tar.add(rootdir, arcname='.')
+    for filename in include:
+        tar.add(filename, arcname=mkrel(filename), recursive=False)
     tar.close()
 
+    include.remove(rootdir)    
+    for filename in reversed(include):
+        if os.path.isdir(filename):
+            os.rmdir(filename)
+        else:
+            os.remove(filename)
+
 
 def unpack_chunk(chunk_filename, dirname):
     '''Unpack a chunk into a directory.
diff --git a/morphlib/bins_tests.py b/morphlib/bins_tests.py
index 64858e7a..02706648 100644
--- a/morphlib/bins_tests.py
+++ b/morphlib/bins_tests.py
@@ -35,7 +35,7 @@ def recursive_lstat(root):
         if pathname == root:
             return '.'
         else:
-            return pathname[len(root):]
+            return pathname[len(root)+1:]
     
     def lstat(filename):
         st = os.lstat(filename)
@@ -79,28 +79,28 @@ class ChunkTests(unittest.TestCase):
 
     def test_empties_everything(self):
         self.populate_instdir()
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['*'])
+        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['.'])
         empty = os.path.join(self.tempdir, 'empty')
         os.mkdir(empty)
-        self.assertEqual(recursive_lstat(self.instdir), empty)
+        self.assertEqual([x for x,y in recursive_lstat(self.instdir)], ['.'])
 
     def test_creates_and_unpacks_chunk_exactly(self):
         self.populate_instdir()
         orig_files = recursive_lstat(self.instdir)
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['*'])
+        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['.'])
         os.mkdir(self.unpacked)
         morphlib.bins.unpack_chunk(self.chunk_file, self.unpacked)
         self.assertEqual(orig_files, recursive_lstat(self.unpacked))
 
     def test_uses_only_matching_names(self):
         self.populate_instdir()
-        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['bin/*'])
+        morphlib.bins.create_chunk(self.instdir, self.chunk_file, ['bin'])
         os.mkdir(self.unpacked)
         morphlib.bins.unpack_chunk(self.chunk_file, self.unpacked)
         self.assertEqual([x for x,y in recursive_lstat(self.unpacked)],
-                         ['bin/foo'])
+                         ['.', 'bin', 'bin/foo'])
         self.assertEqual([x for x,y in recursive_lstat(self.instdir)],
-                         ['lib/libfoo.so'])
+                         ['.', 'lib', 'lib/libfoo.so'])
 
 class StratumTests(unittest.TestCase):