From 521c7e19164fc7eb60a39c4efc5130f7fc814031 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= Date: Tue, 27 Nov 2018 14:20:08 +0100 Subject: [PATCH] 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 | 13 ++++- blender_asset_tracer/compressor.py | 77 +++++++++++++++++++++++++ blender_asset_tracer/pack/__init__.py | 8 ++- blender_asset_tracer/pack/filesystem.py | 21 ++++++- tests/test_compressor.py | 73 +++++++++++++++++++++++ tests/test_pack.py | 35 ++++++++++- tests/test_tracer.py | 2 +- 7 files changed, 221 insertions(+), 8 deletions(-) create mode 100644 blender_asset_tracer/compressor.py create mode 100644 tests/test_compressor.py diff --git a/blender_asset_tracer/cli/pack.py b/blender_asset_tracer/cli/pack.py index b1a3d06..074e35c 100644 --- a/blender_asset_tracer/cli/pack.py +++ b/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. diff --git a/blender_asset_tracer/compressor.py b/blender_asset_tracer/compressor.py new file mode 100644 index 0000000..c800322 --- /dev/null +++ b/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() diff --git a/blender_asset_tracer/pack/__init__.py b/blender_asset_tracer/pack/__init__.py index 1738ef3..3a62b73 100644 --- a/blender_asset_tracer/pack/__init__.py +++ b/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): diff --git a/blender_asset_tracer/pack/filesystem.py b/blender_asset_tracer/pack/filesystem.py index c2daf33..8e5ae06 100644 --- a/blender_asset_tracer/pack/filesystem.py +++ b/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) diff --git a/tests/test_compressor.py b/tests/test_compressor.py new file mode 100644 index 0000000..287bec0 --- /dev/null +++ b/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) diff --git a/tests/test_pack.py b/tests/test_pack.py index d12a3c1..4fed394 100644 --- a/tests/test_pack.py +++ b/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): diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 94f9583..8cca681 100644 --- a/tests/test_tracer.py +++ b/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