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.
This commit is contained in:
Sybren A. Stüvel 2018-11-27 14:20:08 +01:00
parent c1aaa3aab3
commit 521c7e1916
7 changed files with 221 additions and 8 deletions

View File

@ -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.

View File

@ -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()

View File

@ -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):

View File

@ -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)

73
tests/test_compressor.py Normal file
View File

@ -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)

View File

@ -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):

View File

@ -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