diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index f8d3671..1b5a8b3 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/blender_asset_tracer/blendfile/__init__.py @@ -63,6 +63,7 @@ class BlendFile: self._is_modified = False fileobj = path.open(mode, buffering=FILE_BUFFER_SIZE) + fileobj.seek(0, os.SEEK_SET) magic = fileobj.read(len(BLENDFILE_MAGIC)) if magic == BLENDFILE_MAGIC: @@ -170,6 +171,9 @@ class BlendFile: if not self.fileobj: return + if self._is_modified: + log.debug('closing blend file %s after it was modified', self.raw_filepath) + if self._is_modified and self.is_compressed: log.debug("recompressing modified blend file %s", self.raw_filepath) self.fileobj.seek(os.SEEK_SET, 0) @@ -290,7 +294,7 @@ class BlendFile: abspath = relpath.absolute(root) my_log = self.log.getChild('abspath') - my_log.info('Resolved %s relative to %s to %s', relpath, self.filepath, abspath) + my_log.debug('Resolved %s relative to %s to %s', relpath, self.filepath, abspath) return abspath diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index 4dfa869..154e489 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/blender_asset_tracer/blendfile/dna.py @@ -1,12 +1,14 @@ +import logging +import os import typing -import os - -from . import dna_io, header, exceptions +from . import header, exceptions # Either a simple path b'propname', or a tuple (b'parentprop', b'actualprop', arrayindex) FieldPath = typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]] +log = logging.getLogger(__name__) + class Name: """dna.Name is a C-type name stored in the DNA as bytes.""" @@ -83,6 +85,8 @@ class Field: class Struct: """dna.Struct is a C-type structure stored in the DNA.""" + log = log.getChild('Struct') + def __init__(self, dna_type_id: bytes, size: int = None): """ :param dna_type_id: name of the struct in C, like b'AlembicObjectPath'. @@ -268,7 +272,7 @@ class Struct: struct on disk (e.g. the start of the BlendFileBlock containing the data). """ - assert (type(path) == bytes) + assert isinstance(path, bytes), 'path should be bytes, but is %s' % type(path) field, offset = self.field_from_path(file_header.pointer_size, path) @@ -283,6 +287,11 @@ class Struct: fileobj.seek(offset, os.SEEK_CUR) + if self.log.isEnabledFor(logging.DEBUG): + filepos = fileobj.tell() + thing = 'string' if isinstance(value, str) else 'bytes' + self.log.debug('writing %s %r at file offset %d / %x', thing, value, filepos, filepos) + if isinstance(value, str): return endian.write_string(fileobj, value, dna_name.array_size) else: diff --git a/blender_asset_tracer/bpathlib.py b/blender_asset_tracer/bpathlib.py index 0e836e3..30be080 100644 --- a/blender_asset_tracer/bpathlib.py +++ b/blender_asset_tracer/bpathlib.py @@ -21,6 +21,29 @@ class BlendPath(bytes): return bytes.__new__(cls, path.encode('utf-8')) return bytes.__new__(cls, path) + @classmethod + def mkrelative(cls, asset_path: pathlib.Path, bfile_path: pathlib.Path) -> 'BlendPath': + """Construct a BlendPath to the asset relative to the blend file.""" + from collections import deque + + bdir_parts = deque(bfile_path.absolute().parent.parts) + asset_parts = deque(asset_path.absolute().parts) + + # Remove matching initial parts. What is left in bdir_parts represents + # the number of '..' we need. What is left in asset_parts represents + # what we need after the '../../../'. + while bdir_parts: + if bdir_parts[0] != asset_parts[0]: + break + bdir_parts.popleft() + asset_parts.popleft() + + rel_asset = pathlib.Path(*asset_parts) + # TODO(Sybren): should we use sys.getfilesystemencoding() instead? + rel_bytes = str(rel_asset).encode('utf-8') + as_bytes = b'//' + len(bdir_parts) * b'../' + rel_bytes + return cls(as_bytes) + def __str__(self) -> str: """Decodes the path as UTF-8, replacing undecodable bytes. @@ -76,7 +99,7 @@ class BlendPath(bytes): return False - def absolute(self, root: bytes=None) -> 'BlendPath': + def absolute(self, root: bytes = None) -> 'BlendPath': """Determine absolute path. :param root: root directory to compute paths relative to. diff --git a/blender_asset_tracer/cli/pack.py b/blender_asset_tracer/cli/pack.py index 38c9fda..86f2438 100644 --- a/blender_asset_tracer/cli/pack.py +++ b/blender_asset_tracer/cli/pack.py @@ -32,8 +32,8 @@ def add_parser(subparsers): def cli_pack(args): bpath, ppath, tpath = paths_from_cli(args) packer = pack.Packer(bpath, ppath, tpath, args.noop) - packer.investigate() - packer.pack() + packer.strategise() + packer.execute() def paths_from_cli(args) -> (pathlib.Path, pathlib.Path, pathlib.Path): @@ -45,6 +45,9 @@ def paths_from_cli(args) -> (pathlib.Path, pathlib.Path, pathlib.Path): if not bpath.exists(): log.critical('File %s does not exist', bpath) sys.exit(3) + if bpath.is_dir(): + log.critical('%s is a directory, should be a blend file') + sys.exit(3) tpath = args.target if tpath.exists() and not tpath.is_dir(): @@ -60,10 +63,19 @@ def paths_from_cli(args) -> (pathlib.Path, pathlib.Path, pathlib.Path): if not ppath.exists(): log.critical('Project directory %s does not exist', ppath) sys.exit(5) + if not ppath.is_dir(): + log.warning('Project path %s is not a directory; using the parent %s', ppath, ppath.parent) + ppath = ppath.parent + try: bpath.absolute().relative_to(ppath) except ValueError: log.critical('Project directory %s does not contain blend file %s', args.project, bpath.absolute()) sys.exit(5) + + log.info('Blend file to pack: %s', bpath) + log.info('Project path: %s', ppath) + log.info('Pack will be created in: %s', tpath) + return bpath, ppath, tpath diff --git a/blender_asset_tracer/pack/__init__.py b/blender_asset_tracer/pack/__init__.py index 67a941a..0d8b824 100644 --- a/blender_asset_tracer/pack/__init__.py +++ b/blender_asset_tracer/pack/__init__.py @@ -1,14 +1,30 @@ +import collections +import enum import functools import logging import pathlib import shutil +import typing -from blender_asset_tracer import tracer +from blender_asset_tracer import tracer, bpathlib, blendfile from blender_asset_tracer.cli import common +from blender_asset_tracer.tracer import result log = logging.getLogger(__name__) +class PathAction(enum.Enum): + KEEP_PATH = 1 + FIND_NEW_LOCATION = 2 + + +class AssetAction: + def __init__(self): + self.path_action = PathAction.KEEP_PATH + self.usages = [] # which data blocks are referring to this asset + self.new_path = None + + class Packer: def __init__(self, blendfile: pathlib.Path, @@ -26,41 +42,180 @@ class Packer: if noop: log.warning('Running in no-op mode, only showing what will be done.') - def investigate(self): - pass + # Filled by strategise() + self._actions = collections.defaultdict(AssetAction) + self._rewrites = collections.defaultdict(list) + self._packed_paths = {} # from path in project to path in BAT Pack dir. - def pack(self): + def strategise(self): + """Determine what to do with the assets. + + Places an asset into one of these categories: + - Can be copied as-is, nothing smart required. + - Blend files referring to this asset need to be rewritten. + """ + + # The blendfile that we pack is generally not its own dependency, so + # we have to explicitly add it to the _packed_paths. + bfile_path = self.blendfile.absolute() + self._packed_paths[bfile_path] = self.target / bfile_path.relative_to(self.project) + act = self._actions[bfile_path] + act.path_action = PathAction.KEEP_PATH + + new_location_paths = set() for usage in tracer.deps(self.blendfile): - if usage.asset_path.is_absolute(): - raise NotImplementedError('Sorry, cannot handle absolute paths yet: %s' % usage) + # Needing rewriting is not a per-asset thing, but a per-asset-per- + # blendfile thing, since different blendfiles can refer to it in + # different ways (for example with relative and absolute paths). + asset_path = usage.abspath + bfile_path = usage.block.bfile.filepath.absolute() - for assetpath in usage.files(): - self._copy_to_target(assetpath) + path_in_project = self._path_in_project(asset_path) + use_as_is = usage.asset_path.is_blendfile_relative() and path_in_project + needs_rewriting = not use_as_is - log.info('Copied %d files to %s', len(self._already_copied), self.target) + act = self._actions[asset_path] + assert isinstance(act, AssetAction) - def _copy_to_target(self, assetpath: pathlib.Path): + act.usages.append(usage) + if needs_rewriting: + log.info('%s needs rewritten path to %s', bfile_path, usage.asset_path) + act.path_action = PathAction.FIND_NEW_LOCATION + new_location_paths.add(asset_path) + else: + log.info('%s can keep using %s', bfile_path, usage.asset_path) + self._packed_paths[asset_path] = self.target / asset_path.relative_to(self.project) + + self._find_new_paths(new_location_paths) + self._group_rewrites() + + + def _find_new_paths(self, asset_paths: typing.Set[pathlib.Path]): + """Find new locations in the BAT Pack for the given assets.""" + + for path in asset_paths: + act = self._actions[path] + assert isinstance(act, AssetAction) + # Like a join, but ignoring the fact that 'path' is absolute. + act.new_path = pathlib.Path(self.target, '_outside_project', *path.parts[1:]) + self._packed_paths[path] = act.new_path + + def _group_rewrites(self): + """For each blend file, collect which fields need rewriting. + + This ensures that the execute() step has to visit each blend file + only once. + """ + + for action in self._actions.values(): + if action.path_action != PathAction.FIND_NEW_LOCATION: + # This asset doesn't require a new location, so no rewriting necessary. + continue + + for usage in action.usages: + bfile_path = usage.block.bfile.filepath.absolute().resolve() + self._rewrites[bfile_path].append(usage) + + def _path_in_project(self, path: pathlib.Path) -> bool: try: - assetpath = assetpath.resolve() - except FileNotFoundError: - log.error('Dependency %s does not exist', assetpath) + # MUST use resolve(), otherwise /path/to/proj/../../asset.png + # will return True (relative_to will return ../../asset.png). + path.resolve().relative_to(self.project) + except ValueError: + return False + return True - if assetpath in self._already_copied: - log.debug('Already copied %s', assetpath) - return - self._already_copied.add(assetpath) + def execute(self): + """Execute the strategy.""" + assert self._actions, 'Run strategise() first' - relpath = self._shorten(assetpath) - if relpath.is_absolute(): - raise NotImplementedError( - 'Sorry, cannot handle paths outside project directory yet: %s is not in %s' - % (relpath, self.project)) + self._copy_files_to_target() + if not self.noop: + self._rewrite_paths() + + def _copy_files_to_target(self): + """Copy all assets to the target directoy. + + This creates the BAT Pack but does not yet do any path rewriting. + """ + log.info('Executing %d copy actions', len(self._actions)) + for asset_path, action in self._actions.items(): + self._copy_asset_and_deps(asset_path, action) - full_target = self.target / relpath - full_target.parent.mkdir(parents=True, exist_ok=True) if self.noop: - print('%s → %s' % (assetpath, full_target)) + msg = 'Would copy' else: - print(relpath) - # TODO(Sybren): when we target Py 3.6+, remove the str() calls. - shutil.copyfile(str(assetpath), str(full_target)) + msg = 'Copied' + log.info('%s %d files to %s', msg, len(self._already_copied), self.target) + + def _rewrite_paths(self): + """Rewrite paths to the new location of the assets.""" + + for bfile_path, rewrites in self._rewrites.items(): + assert isinstance(bfile_path, pathlib.Path) + bfile_pp = self._packed_paths[bfile_path] + + log.info('Rewriting %s', bfile_pp) + + with blendfile.BlendFile(bfile_pp, 'rb+') as bfile: + for usage in rewrites: + assert isinstance(usage, result.BlockUsage) + asset_pp = self._packed_paths[usage.abspath] + assert isinstance(asset_pp, pathlib.Path) + + log.debug(' - %s is packed at %s', usage.asset_path, asset_pp) + relpath = bpathlib.BlendPath.mkrelative(asset_pp, bfile_pp) + if relpath == usage.asset_path: + log.info(' - %s remained at %s', usage.asset_path, relpath) + continue + + log.info(' - %s moved to %s', usage.asset_path, relpath) + + # Find the same block in the newly copied file. + block = bfile.dereference_pointer(usage.block.addr_old) + log.info(' - 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) + bfile.fileobj.flush() + + def _copy_asset_and_deps(self, asset_path: pathlib.Path, action: AssetAction): + log.info('Copying %s and dependencies', asset_path) + + # Copy the asset itself. + packed_path = self._packed_paths[asset_path] + self._copy_to_target(asset_path, packed_path) + + # Copy its dependencies. + for usage in action.usages: + for file_path in usage.files(): + # TODO(Sybren): handle sequences properly! + packed_path = self._packed_paths[file_path] + self._copy_to_target(file_path, packed_path) + + def _copy_to_target(self, asset_path: pathlib.Path, target: pathlib.Path): + if self._is_already_copied(asset_path): + return + + print('%s → %s' % (asset_path, target)) + if self.noop: + return + + target.parent.mkdir(parents=True, exist_ok=True) + # TODO(Sybren): when we target Py 3.6+, remove the str() calls. + shutil.copyfile(str(asset_path), str(target)) + + def _is_already_copied(self, asset_path: pathlib.Path) -> bool: + try: + asset_path = asset_path.resolve() + except FileNotFoundError: + log.error('Dependency %s does not exist', asset_path) + return True + + if asset_path in self._already_copied: + log.debug('Already copied %s', asset_path) + return True + + # Assume the copy will happen soon. + self._already_copied.add(asset_path) + return False diff --git a/blender_asset_tracer/tracer/result.py b/blender_asset_tracer/tracer/result.py index 9c2c2d6..e26da08 100644 --- a/blender_asset_tracer/tracer/result.py +++ b/blender_asset_tracer/tracer/result.py @@ -67,6 +67,7 @@ class BlockUsage: self.path_full_field = path_full_field self.path_dir_field = path_dir_field self.path_base_field = path_base_field + self._abspath = None # cached by __fspath__() @staticmethod def guess_block_name(block: blendfile.BlendFileBlock) -> bytes: @@ -97,8 +98,7 @@ class BlockUsage: It is assumed that paths are valid UTF-8. """ - bpath = self.block.bfile.abspath(self.asset_path) - path = bpath.to_path() + path = self.__fspath__() if not self.is_sequence: if not path.exists(): log.warning('Path %s does not exist for %s', path, self) @@ -110,3 +110,12 @@ class BlockUsage: yield from file_sequence.expand_sequence(path) except file_sequence.DoesNotExist: log.warning('Path %s does not exist for %s', path, self) + + def __fspath__(self) -> pathlib.Path: + """Determine the absolute path of the asset on the filesystem.""" + if self._abspath is None: + bpath = self.block.bfile.abspath(self.asset_path) + self._abspath = bpath.to_path().resolve() + return self._abspath + + abspath = property(__fspath__) diff --git a/tests/blendfiles/subdir/doubly_linked_up.blend b/tests/blendfiles/subdir/doubly_linked_up.blend new file mode 100644 index 0000000..fe45cb4 Binary files /dev/null and b/tests/blendfiles/subdir/doubly_linked_up.blend differ diff --git a/tests/test_bpathlib.py b/tests/test_bpathlib.py index c91c83b..b6fcb4f 100644 --- a/tests/test_bpathlib.py +++ b/tests/test_bpathlib.py @@ -1,3 +1,4 @@ +from pathlib import Path import unittest from blender_asset_tracer.bpathlib import BlendPath @@ -52,3 +53,29 @@ class BlendPathTest(unittest.TestCase): self.assertEqual(b'/root/and/parent.blend', b'/root/and' / BlendPath(b'parent.blend')) with self.assertRaises(ValueError): b'/root/and' / BlendPath(b'/parent.blend') + + def test_mkrelative(self): + self.assertEqual(b'//asset.png', BlendPath.mkrelative( + Path('/path/to/asset.png'), + Path('/path/to/bfile.blend'), + )) + self.assertEqual(b'//to/asset.png', BlendPath.mkrelative( + Path('/path/to/asset.png'), + Path('/path/bfile.blend'), + )) + self.assertEqual(b'//../of/asset.png', BlendPath.mkrelative( + Path('/path/of/asset.png'), + Path('/path/to/bfile.blend'), + )) + self.assertEqual(b'//../../path/of/asset.png', BlendPath.mkrelative( + Path('/path/of/asset.png'), + Path('/some/weird/bfile.blend'), + )) + self.assertEqual(b'//very/very/very/very/very/deep/asset.png', BlendPath.mkrelative( + Path('/path/to/very/very/very/very/very/deep/asset.png'), + Path('/path/to/bfile.blend'), + )) + self.assertEqual(b'//../../../../../../../../shallow/asset.png', BlendPath.mkrelative( + Path('/shallow/asset.png'), + Path('/path/to/very/very/very/very/very/deep/bfile.blend'), + )) diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 5e05da2..e32394a 100644 --- a/tests/test_tracer.py +++ b/tests/test_tracer.py @@ -212,6 +212,14 @@ class DepsTest(AbstractTracerTest): # b'//textures/Textures/Buildings/buildings_roof_04-color.png', False), }) + def test_usage_abspath(self): + deps = [dep for dep in tracer.deps(self.blendfiles / 'doubly_linked.blend') + if dep.asset_path == b'//material_textures.blend'] + usage = deps[0] + + expect = self.blendfiles / 'material_textures.blend' + self.assertEqual(expect, usage.abspath) + def test_sim_data(self): self.assert_deps('T53562/bam_pack_bug.blend', { b'OBEmitter.modifiers[0]': Expect(