Added --compress option for 'bat pack' command

This compresses all packed Blend files. Other files, as well as already-
compressed Blend files, are left as-is.

blender_asset_tracer/cli/pack.py

@@ -47,6 +47,8 @@ def add_parser(subparsers):
                         help="Don't copy files, just show what would be done.")
     parser.add_argument('-e', '--exclude', nargs='*', default='',
                         help="Space-separated list of glob patterns (like '*.abc') to exclude.")
+    parser.add_argument('-c', '--compress', default=False, action='store_true',
+                        help='Compress blend files while copying')
 
 
 def cli_pack(args):
@@ -68,12 +70,19 @@ def create_packer(args, bpath: pathlib.Path, ppath: pathlib.Path,
         if args.noop:
             raise ValueError('S3 uploader does not support no-op.')
 
+        if args.compress:
+            raise ValueError('S3 uploader does not support on-the-fly compression')
+
         packer = create_s3packer(bpath, ppath, tpath)
     elif tpath.suffix.lower() == '.zip':
         from blender_asset_tracer.pack import zipped
-        packer = zipped.ZipPacker(bpath, ppath, tpath, args.noop)
+
+        if args.compress:
+            raise ValueError('ZIP packer does not support on-the-fly compression')
+
+        packer = zipped.ZipPacker(bpath, ppath, tpath, noop=args.noop)
     else:
-        packer = pack.Packer(bpath, ppath, tpath, args.noop)
+        packer = pack.Packer(bpath, ppath, tpath, noop=args.noop, compress=args.compress)
 
     if args.exclude:
         # args.exclude is a list, due to nargs='*', so we have to split and flatten.

blender_asset_tracer/compressor.py

@@ -0,0 +1,77 @@
+"""shutil-like functionality while compressing blendfiles on the fly."""
+
+import gzip
+import logging
+import pathlib
+import shutil
+
+log = logging.getLogger(__name__)
+
+# Arbitrarily chosen block size, in bytes.
+BLOCK_SIZE = 256 * 2 ** 10
+
+
+def move(src: pathlib.Path, dest: pathlib.Path):
+    """Move a file from src to dest, gzip-compressing if not compressed yet.
+
+    Only compresses files ending in .blend; others are moved as-is.
+    """
+    my_log = log.getChild('move')
+    my_log.debug('Moving %s to %s', src, dest)
+
+    if src.suffix.lower() == '.blend':
+        _move_or_copy(src, dest, my_log, source_must_remain=False)
+    else:
+        shutil.move(str(src), str(dest))
+
+
+def copy(src: pathlib.Path, dest: pathlib.Path):
+    """Copy a file from src to dest, gzip-compressing if not compressed yet.
+
+    Only compresses files ending in .blend; others are copied as-is.
+    """
+    my_log = log.getChild('copy')
+    my_log.debug('Copying %s to %s', src, dest)
+
+    if src.suffix.lower() == '.blend':
+        _move_or_copy(src, dest, my_log, source_must_remain=True)
+    else:
+        shutil.copy2(str(src), str(dest))
+
+
+def _move_or_copy(src: pathlib.Path, dest: pathlib.Path,
+                  my_log: logging.Logger,
+                  *,
+                  source_must_remain: bool):
+    """Either move or copy a file, gzip-compressing if not compressed yet.
+
+    :param src: File to copy/move.
+    :param dest: Path to copy/move to.
+    :source_must_remain: True to copy, False to move.
+    :my_log: Logger to use for logging.
+    """
+    srcfile = src.open('rb')
+    try:
+        first_bytes = srcfile.read(2)
+        if first_bytes == b'\x1f\x8b':
+            # Already a gzipped file.
+            srcfile.close()
+            my_log.debug('Source file %s is GZipped already', src)
+            if source_must_remain:
+                shutil.copy2(str(src), str(dest))
+            else:
+                shutil.move(str(src), str(dest))
+            return
+
+        my_log.debug('Compressing %s on the fly while copying to %s', src, dest)
+        with gzip.open(str(dest), mode='wb') as destfile:
+            destfile.write(first_bytes)
+            shutil.copyfileobj(srcfile, destfile, BLOCK_SIZE)
+
+        srcfile.close()
+        if not source_must_remain:
+            my_log.debug('Deleting source file %s', src)
+            src.unlink()
+    finally:
+        if not srcfile.closed:
+            srcfile.close()

blender_asset_tracer/pack/__init__.py

@@ -97,11 +97,14 @@ class Packer:
                  bfile: pathlib.Path,
                  project: pathlib.Path,
                  target: pathlib.Path,
-                 noop=False) -> None:
+                 *,
+                 noop=False,
+                 compress=False) -> None:
         self.blendfile = bfile
         self.project = project
         self.target = target
         self.noop = noop
+        self.compress = compress
         self._aborted = threading.Event()
         self._abort_lock = threading.RLock()
 
@@ -352,6 +355,9 @@ class Packer:
 
     def _create_file_transferer(self) -> transfer.FileTransferer:
         """Create a FileCopier(), can be overridden in a subclass."""
+
+        if self.compress:
+            return filesystem.CompressedFileCopier()
         return filesystem.FileCopier()
 
     def _start_file_transferrer(self):

blender_asset_tracer/pack/filesystem.py

@@ -22,6 +22,7 @@ import pathlib
 import shutil
 import typing
 
+from .. import compressor
 from . import transfer
 
 log = logging.getLogger(__name__)
@@ -81,9 +82,17 @@ class FileCopier(transfer.FileTransferer):
         if self.files_skipped:
             log.info('Skipped %d files', self.files_skipped)
 
+    def _move(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
+        """Low-level file move"""
+        shutil.move(str(srcpath), str(dstpath))
+
+    def _copy(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
+        """Low-level file copy"""
+        shutil.copy2(str(srcpath), str(dstpath))
+
     def move(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
         s_stat = srcpath.stat()
-        shutil.move(str(srcpath), str(dstpath))
+        self._move(srcpath, dstpath)
 
         self.files_transferred += 1
         self.report_transferred(s_stat.st_size)
@@ -108,7 +117,7 @@ class FileCopier(transfer.FileTransferer):
                 return
 
         log.debug('Copying %s → %s', srcpath, dstpath)
-        shutil.copy2(str(srcpath), str(dstpath))
+        self._copy(srcpath, dstpath)
 
         self.already_copied.add((srcpath, dstpath))
         self.files_transferred += 1
@@ -176,3 +185,11 @@ class FileCopier(transfer.FileTransferer):
         self.already_copied.add((src, dst))
 
         return dst
+
+
+class CompressedFileCopier(FileCopier):
+    def _move(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
+        compressor.move(srcpath, dstpath)
+
+    def _copy(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
+        compressor.copy(srcpath, dstpath)

tests/test_compressor.py

@@ -0,0 +1,73 @@
+import pathlib
+import tempfile
+import shutil
+
+from blender_asset_tracer import blendfile
+from abstract_test import AbstractBlendFileTest
+
+from blender_asset_tracer import compressor
+
+
+class CompressorTest(AbstractBlendFileTest):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        tempdir = pathlib.Path(self.temp.name)
+        self.srcdir = tempdir / 'src'
+        self.destdir = tempdir / 'dest'
+
+        self.srcdir.mkdir()
+        self.destdir.mkdir()
+
+    def tearDown(self):
+        self.temp.cleanup()
+
+    def _test(self, filename: str, source_must_remain: bool):
+        """Do a move/copy test.
+
+        The result should be the same, regardless of whether the
+        source file was already compressed or not.
+        """
+        # Make a copy we can move around without moving the actual file in
+        # the source tree.
+        srcfile = self.srcdir / filename
+        destfile = self.destdir / filename
+        srcfile.parent.mkdir(parents=True, exist_ok=True)
+        destfile.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(str(self.blendfiles / filename), str(srcfile))
+
+        if source_must_remain:
+            compressor.copy(srcfile, destfile)
+        else:
+            compressor.move(srcfile, destfile)
+
+        self.assertEqual(source_must_remain, srcfile.exists())
+        self.assertTrue(destfile.exists())
+
+        if destfile.suffix == '.blend':
+            self.bf = blendfile.BlendFile(destfile)
+            self.assertTrue(self.bf.is_compressed)
+            return
+
+        with destfile.open('rb') as infile:
+            magic = infile.read(3)
+        if destfile.suffix == '.jpg':
+            self.assertEqual(b'\xFF\xD8\xFF', magic,
+                             'Expected %s to be a JPEG' % destfile)
+        else:
+            self.assertNotEqual(b'\x1f\x8b', magic[:2],
+                                'Expected %s to be NOT compressed' % destfile)
+
+    def test_move_already_compressed(self):
+        self._test('basic_file_ñønæščii.blend', False)
+
+    def test_move_compress_on_the_fly(self):
+        self._test('basic_file.blend', False)
+
+    def test_copy_already_compressed(self):
+        self._test('basic_file_ñønæščii.blend', True)
+
+    def test_copy_compress_on_the_fly(self):
+        self._test('basic_file.blend', True)
+
+    def test_move_jpeg(self):
+        self._test('textures/Bricks/brick_dotted_04-color.jpg', False)

tests/test_pack.py

@@ -14,6 +14,7 @@ class AbstractPackTest(AbstractBlendFileTest):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
+        logging.getLogger('blender_asset_tracer.compressor').setLevel(logging.DEBUG)
         logging.getLogger('blender_asset_tracer.pack').setLevel(logging.DEBUG)
         logging.getLogger('blender_asset_tracer.blendfile.open_cached').setLevel(logging.DEBUG)
         logging.getLogger('blender_asset_tracer.blendfile.open_cached').setLevel(logging.DEBUG)
@@ -23,8 +24,6 @@ class AbstractPackTest(AbstractBlendFileTest):
         super().setUp()
         self.tdir = tempfile.TemporaryDirectory(suffix='-packtest')
         self.tpath = pathlib.Path(self.tdir.name)
-        # self.tpath = pathlib.Path('/tmp/tempdir-packtest')
-        # self.tpath.mkdir(parents=True, exist_ok=True)
 
     def tearDown(self):
         self.tdir.cleanup()
@@ -268,6 +267,38 @@ class PackTest(AbstractPackTest):
         info = infopath.open().read().splitlines(keepends=False)
         self.assertEqual(blendname, info[-1].strip())
 
+    def test_compression(self):
+        blendname = 'subdir/doubly_linked_up.blend'
+        imgfile = self.blendfiles / blendname
+
+        packer = pack.Packer(imgfile, self.blendfiles, self.tpath, compress=True)
+        packer.strategise()
+        packer.execute()
+
+        dest = self.tpath / blendname
+        self.assertTrue(dest.exists())
+        self.assertTrue(blendfile.open_cached(dest).is_compressed)
+
+        for bpath in self.tpath.rglob('*.blend'):
+            if bpath == dest:
+                # Only test files that were bundled as dependency; the main
+                # file was tested above already.
+                continue
+            self.assertTrue(blendfile.open_cached(bpath).is_compressed,
+                            'Expected %s to be compressed' % bpath)
+            break
+        else:
+            self.fail('Expected to have Blend files in the BAT pack.')
+
+        for imgpath in self.tpath.rglob('*.jpg'):
+            with imgpath.open('rb') as imgfile:
+                magic = imgfile.read(3)
+            self.assertEqual(b'\xFF\xD8\xFF', magic,
+                             'Expected %s to NOT be compressed' % imgpath)
+            break
+        else:
+            self.fail('Expected to have JPEG files in the BAT pack.')
+
 
 class ProgressTest(AbstractPackTest):
     def test_strategise(self):

tests/test_tracer.py

@@ -253,7 +253,7 @@ class DepsTest(AbstractTracerTest):
 
         reclim = sys.getrecursionlimit()
         try:
-            sys.setrecursionlimit(80)
+            sys.setrecursionlimit(100)
             # This should finish without hitting the recursion limit.
             for _ in trace.deps(infinite_bfile):
                 pass