autour_de_minuit/blender-asset-tracer
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

Rewrite blend files in temporary directory, before copying

Browse Source 
This allows us to later support uploading to a non-local filesystem, such
as Amazon S3 cloud storage, after path rewriting is complete.
This commit is contained in:
Sybren A. Stüvel 2018-03-09 12:33:15 +01:00
parent 384f1d5107
commit 824ca4bcb7
4 changed files with 212 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
blender_asset_tracer/blendfile/__init__.py
View File

@ -28,6 +28,7 @@ import logging
import os
import struct
import pathlib
import shutil
import tempfile
import typing
@ -44,19 +45,28 @@ BFBList = typing.List['BlendFileBlock']
_cached_bfiles = {} # type: typing.Dict[pathlib.Path, BlendFile]
def open_cached(path: pathlib.Path, mode='rb', assert_cached=False) -> 'BlendFile':
def open_cached(path: pathlib.Path, mode='rb',
assert_cached: typing.Optional[bool] = None) -> 'BlendFile':
"""Open a blend file, ensuring it is only opened once."""
my_log = log.getChild('open_cached')
bfile_path = path.absolute().resolve()
try:
return _cached_bfiles[bfile_path]
except KeyError:
pass
if assert_cached:
if assert_cached is not None:
is_cached = bfile_path in _cached_bfiles
if assert_cached and not is_cached:
raise AssertionError('File %s was not cached' % bfile_path)
elif not assert_cached and is_cached:
raise AssertionError('File %s was cached' % bfile_path)
try:
bfile = _cached_bfiles[bfile_path]
except KeyError:
my_log.debug('Opening non-cached %s', path)
bfile = BlendFile(path, mode=mode)
_cached_bfiles[bfile_path] = bfile
else:
my_log.debug('Returning cached %s', path)
return bfile
@ -67,7 +77,7 @@ def close_all_cached():
return
log.debug('Closing %d cached blend files', len(_cached_bfiles))
for bfile in _cached_bfiles.values():
for bfile in list(_cached_bfiles.values()):
bfile.close()
_cached_bfiles.clear()
@ -136,6 +146,8 @@ class BlendFile:
if 'b' not in mode:
raise ValueError('Only binary modes are supported, not %r' % mode)
self.filepath = path
fileobj = path.open(mode, buffering=FILE_BUFFER_SIZE) # typing.IO[bytes]
fileobj.seek(0, os.SEEK_SET)
@ -206,7 +218,7 @@ class BlendFile:
def __exit__(self, exctype, excvalue, traceback):
self.close()
def rebind(self, path: pathlib.Path, mode='rb'):
def copy_and_rebind(self, path: pathlib.Path, mode='rb'):
"""Change which file is bound to this BlendFile.
This allows cloning a previously opened file, and rebinding it to reuse
@ -217,6 +229,10 @@ class BlendFile:
self.close()
_uncache(self.filepath)
self.log.debug('Copying %s to %s', self.filepath, path)
# TODO(Sybren): remove str() calls when targeting Python 3.6+
shutil.copy(str(self.filepath), str(path))
self._open_file(path, mode=mode)
_cache(path, self)
@ -261,6 +277,11 @@ class BlendFile:
self.fileobj.close()
self._is_modified = False
try:
del _cached_bfiles[self.filepath]
except KeyError:
pass
def ensure_subtype_smaller(self, sdna_index_curr, sdna_index_next):
# never refine to a smaller type
curr_struct = self.structs[sdna_index_curr]

68
blender_asset_tracer/pack/__init__.py
View File

@ -3,6 +3,8 @@ import enum
import functools
import logging
import pathlib
import tempfile
import typing
from blender_asset_tracer import trace, bpathlib, blendfile
@ -33,7 +35,12 @@ class AssetAction:
self.new_path = None # type: pathlib.Path
"""Absolute path to the asset in the BAT Pack."""
self.read_from = None # type: pathlib.Path
self.read_from = None # type: typing.Optional[pathlib.Path]
"""Optional path from which to read the asset.
This is used when blend files have been rewritten. It is assumed that
when this property is set, the file can be moved instead of copied.
"""
self.rewrites = [] # type: typing.List[result.BlockUsage]
"""BlockUsage objects in this asset that may require rewriting.
@ -65,6 +72,19 @@ class Packer:
# Number of files we would copy, if not for --noop
self._file_count = 0
self._tmpdir = tempfile.TemporaryDirectory(suffix='-batpack')
self._rewrite_in = pathlib.Path(self._tmpdir.name)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
def close(self):
"""Clean up any temporary files."""
self._tmpdir.cleanup()
def strategise(self):
"""Determine what to do with the assets.
@ -148,9 +168,9 @@ class Packer:
"""Execute the strategy."""
assert self._actions, 'Run strategise() first'
self._copy_files_to_target()
if not self.noop:
self._rewrite_paths()
self._copy_files_to_target()
def _copy_files_to_target(self):
"""Copy all assets to the target directoy.
@ -172,21 +192,36 @@ class Packer:
fc.done_and_join()
def _rewrite_paths(self):
"""Rewrite paths to the new location of the assets."""
"""Rewrite paths to the new location of the assets.
Writes the rewritten blend files to a temporary location.
"""
for bfile_path, action in self._actions.items():
if not action.rewrites:
continue
assert isinstance(bfile_path, pathlib.Path)
# bfile_pp is the final path of this blend file in the BAT pack.
# It is used to determine relative paths to other blend files.
# It is *not* used for any disk I/O, since the file may not even
# exist on the local filesystem.
bfile_pp = self._actions[bfile_path].new_path
log.info('Rewriting %s', bfile_pp)
# Use tempfile to create a unique name in our temporary directoy.
# The file should be deleted when self.close() is called, and not
# when the bfile_tp object is GC'd.
bfile_tmp = tempfile.NamedTemporaryFile(dir=str(self._rewrite_in),
suffix='-' + bfile_path.name,
delete=False)
bfile_tp = pathlib.Path(bfile_tmp.name)
action.read_from = bfile_tp
log.info('Rewriting %s to %s', bfile_path, bfile_tp)
# The original blend file will have been cached, so we can use it
# to avoid re-parsing all data blocks in the to-be-rewritten file.
bfile = blendfile.open_cached(bfile_path, assert_cached=True)
bfile.rebind(bfile_pp, mode='rb+')
bfile.copy_and_rebind(bfile_tp, mode='rb+')
for usage in action.rewrites:
assert isinstance(usage, result.BlockUsage)
@ -204,11 +239,11 @@ class Packer:
# Find the same block in the newly copied file.
block = bfile.dereference_pointer(usage.block.addr_old)
if usage.path_full_field is None:
log.info(' - updating field %s of block %s',
log.debug(' - updating field %s of block %s',
usage.path_dir_field.name.name_only, block)
reldir = bpathlib.BlendPath.mkrelative(asset_pp.parent, bfile_pp)
written = block.set(usage.path_dir_field.name.name_only, reldir)
log.info(' - written %d bytes', written)
log.debug(' - written %d bytes', written)
# BIG FAT ASSUMPTION that the filename (e.g. basename
# without path) does not change. This makes things much
@ -216,18 +251,24 @@ class Packer:
# filename fields are in different blocks. See the
# blocks2assets.scene() function for the implementation.
else:
log.info(' - updating field %s of block %s',
log.debug(' - updating field %s of block %s',
usage.path_full_field.name.name_only, block)
written = block.set(usage.path_full_field.name.name_only, relpath)
log.info(' - written %d bytes', written)
log.debug(' - written %d bytes', written)
# Make sure we close the file, otherwise changes may not be
# flushed before it gets copied.
bfile.close()
def _copy_asset_and_deps(self, asset_path: pathlib.Path, action: AssetAction,
fc: queued_copy.FileCopier):
log.debug('Queueing copy of %s and dependencies', asset_path)
# Copy the asset itself.
packed_path = self._actions[asset_path].new_path
self._copy_to_target(asset_path, packed_path, fc)
packed_path = action.new_path
read_path = action.read_from or asset_path
# TODO(Sybren): if the asset is a rewritten blend file (and thus a copy),
# do a move instead of a copy.
self._copy_to_target(read_path, packed_path, fc)
# Copy its sequence dependencies.
for usage in action.usages:
@ -253,4 +294,5 @@ class Packer:
print('%s → %s' % (asset_path, target))
self._file_count += 1
return
log.debug('Queueing copy of %s', asset_path)
fc.queue(asset_path, target)

91
tests/test_blendfile_loading.py
View File

@ -1,5 +1,6 @@
import os
import pathlib
import tempfile
from blender_asset_tracer import blendfile
from blender_asset_tracer.blendfile import iterators, exceptions
@ -284,3 +285,93 @@ class LoadNonBlendfileTest(AbstractBlendFileTest):
def test_no_datablocks(self):
with self.assertRaises(exceptions.NoDNA1Block):
blendfile.BlendFile(self.blendfiles / 'corrupt_only_magic.blend')
class BlendFileCacheTest(AbstractBlendFileTest):
def setUp(self):
super().setUp()
self.tdir = tempfile.TemporaryDirectory()
self.tpath = pathlib.Path(self.tdir.name)
def tearDown(self):
self.tdir.cleanup()
def test_open_cached(self):
infile = self.blendfiles / 'basic_file.blend'
bf1 = blendfile.open_cached(infile)
bf2 = blendfile.open_cached(infile)
# The file should only be opened & parsed once.
self.assertIs(bf1, bf2)
self.assertIs(bf1, blendfile._cached_bfiles[infile])
def test_compressed(self):
infile = self.blendfiles / 'linked_cube_compressed.blend'
bf1 = blendfile.open_cached(infile)
bf2 = blendfile.open_cached(infile)
# The file should only be opened & parsed once.
self.assertIs(bf1, bf2)
self.assertIs(bf1, blendfile._cached_bfiles[infile])
def test_closed(self):
infile = self.blendfiles / 'linked_cube_compressed.blend'
bf = blendfile.open_cached(infile)
self.assertIs(bf, blendfile._cached_bfiles[infile])
blendfile.close_all_cached()
self.assertTrue(bf.fileobj.closed)
self.assertEqual({}, blendfile._cached_bfiles)
def test_close_one_file(self):
path1 = self.blendfiles / 'linked_cube_compressed.blend'
path2 = self.blendfiles / 'basic_file.blend'
bf1 = blendfile.open_cached(path1)
bf2 = blendfile.open_cached(path2)
self.assertIs(bf1, blendfile._cached_bfiles[path1])
# Closing a file should remove it from the cache.
bf1.close()
self.assertTrue(bf1.fileobj.closed)
self.assertEqual({path2: bf2}, blendfile._cached_bfiles)
def test_open_and_rebind(self):
infile = self.blendfiles / 'linked_cube.blend'
other = self.tpath / 'copy.blend'
self._open_and_rebind_test(infile, other)
def test_open_and_rebind_compressed(self):
infile = self.blendfiles / 'linked_cube_compressed.blend'
other = self.tpath / 'copy.blend'
self._open_and_rebind_test(infile, other)
def _open_and_rebind_test(self, infile: pathlib.Path, other: pathlib.Path):
self.assertFalse(other.exists())
bf = blendfile.open_cached(infile)
self.assertEqual(str(bf.raw_filepath), bf.fileobj.name)
before_filepath = bf.filepath
before_raw_fp = bf.raw_filepath
before_blocks = bf.blocks
before_compressed = bf.is_compressed
bf.copy_and_rebind(other, mode='rb+')
self.assertTrue(other.exists())
self.assertEqual(before_compressed, bf.is_compressed)
if bf.is_compressed:
self.assertNotEqual(bf.filepath, bf.raw_filepath)
else:
self.assertEqual(bf.filepath, bf.raw_filepath)
self.assertNotEqual(before_filepath, bf.filepath)
self.assertNotEqual(before_raw_fp, bf.raw_filepath)
self.assertEqual(other, bf.filepath)
self.assertIs(before_blocks, bf.blocks)
self.assertNotIn(infile, blendfile._cached_bfiles)
self.assertIs(bf, blendfile._cached_bfiles[other])
self.assertEqual(str(bf.raw_filepath), bf.fileobj.name)

53
tests/test_pack.py
View File

@ -15,11 +15,16 @@ class AbstractPackTest(AbstractBlendFileTest):
def setUpClass(cls):
super().setUpClass()
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)
logging.getLogger('blender_asset_tracer.blendfile.BlendFile').setLevel(logging.DEBUG)
def setUp(self):
super().setUp()
self.tdir = tempfile.TemporaryDirectory()
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()
@ -112,7 +117,33 @@ class AbstractPackTest(AbstractBlendFileTest):
self.assertEqual(b'LILib.002', rw_dbllink[1].block_name)
self.assertEqual(b'//../material_textures.blend', rw_dbllink[1].asset_path)
def test_execute_rewrite_no_touch_origs(self):
infile = self._pack_with_rewrite()
# The original file shouldn't be touched.
bfile = blendfile.open_cached(infile, assert_cached=False)
libs = sorted(bfile.code_index[b'LI'])
self.assertEqual(b'LILib', libs[0].id_name)
self.assertEqual(b'//../linked_cube.blend', libs[0][b'name'])
self.assertEqual(b'LILib.002', libs[1].id_name)
self.assertEqual(b'//../material_textures.blend', libs[1][b'name'])
def test_execute_rewrite(self):
infile = self._pack_with_rewrite()
extpath = pathlib.Path('//_outside_project', *self.blendfiles.parts[1:])
extbpath = bpathlib.BlendPath(extpath)
# Those libraries should be properly rewritten.
bfile = blendfile.open_cached(self.tpath / infile.name, assert_cached=False)
libs = sorted(bfile.code_index[b'LI'])
self.assertEqual(b'LILib', libs[0].id_name)
self.assertEqual(extbpath / b'linked_cube.blend', libs[0][b'name'])
self.assertEqual(b'LILib.002', libs[1].id_name)
self.assertEqual(extbpath / b'material_textures.blend', libs[1][b'name'])
def _pack_with_rewrite(self):
ppath = self.blendfiles / 'subdir'
infile = ppath / 'doubly_linked_up.blend'
@ -120,25 +151,7 @@ class AbstractPackTest(AbstractBlendFileTest):
packer.strategise()
packer.execute()
extpath = pathlib.Path('//_outside_project', *self.blendfiles.parts[1:])
extbpath = bpathlib.BlendPath(extpath)
# Those libraries should be properly rewritten.
bfile = blendfile.open_cached(self.tpath / infile.name)
libs = sorted(bfile.code_index[b'LI'])
self.assertEqual(b'LILib', libs[0].id_name)
self.assertEqual(extbpath / b'linked_cube.blend', libs[0][b'name'])
self.assertEqual(b'LILib.002', libs[1].id_name)
self.assertEqual(extbpath / b'material_textures.blend', libs[1][b'name'])
# The original file shouldn't be touched, though.
bfile = blendfile.open_cached(infile)
libs = sorted(bfile.code_index[b'LI'])
self.assertEqual(b'LILib', libs[0].id_name)
self.assertEqual(b'//../linked_cube.blend', libs[0][b'name'])
self.assertEqual(b'LILib.002', libs[1].id_name)
self.assertEqual(b'//../material_textures.blend', libs[1][b'name'])
return infile
def test_noop(self):
ppath = self.blendfiles / 'subdir'