blender_asset_tracer/pack/__init__.py

# ***** BEGIN GPL LICENSE BLOCK *****
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
# ***** END GPL LICENCE BLOCK *****
#
# (c) 2018, Blender Foundation - Sybren A. Stüvel
import collections
import enum
import functools
import logging
import pathlib
import tempfile
import threading
import typing

from blender_asset_tracer import trace, bpathlib, blendfile
from blender_asset_tracer.trace import file_sequence, result
from . import filesystem, transfer, progress

log = logging.getLogger(__name__)


class PathAction(enum.Enum):
    KEEP_PATH = 1
    FIND_NEW_LOCATION = 2


class AssetAction:
    """All the info required to rewrite blend files and copy assets."""

    def __init__(self) -> None:
        self.path_action = PathAction.KEEP_PATH
        self.usages = []  # type: typing.List[result.BlockUsage]
        """BlockUsage objects referring to this asset.

        Those BlockUsage objects could refer to data blocks in this blend file
        (if the asset is a blend file) or in another blend file.
        """

        self.new_path = None  # type: typing.Optional[pathlib.PurePath]
        """Absolute path to the asset in the BAT Pack.

        This path may not exist on the local file system at all, for example
        when uploading files to remote S3-compatible storage.
        """

        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.

        Empty list if this AssetAction is not for a blend file.
        """


class Aborted(RuntimeError):
    """Raised by Packer to abort the packing process.

    See the Packer.abort() function.
    """


class Packer:
    """Takes a blend file and bundle it with its dependencies.

    The process is separated into two functions:

        - strategise() finds all the dependencies and determines what to do
          with them.
        - execute() performs the actual packing operation, by rewriting blend
          files to ensure the paths to moved files are correct and
          transferring the files.

    The file transfer is performed in a separate thread by a FileTransferer
    instance.
    """

    def __init__(self,
                 bfile: pathlib.Path,
                 project: pathlib.Path,
                 target: str,
                 *,
                 noop=False,
                 compress=False,
                 relative_only=False) -> None:
        self.blendfile = bfile
        self.project = project
        self.target = target
        self._target_path = self._make_target_path(target)
        self.noop = noop
        self.compress = compress
        self.relative_only = relative_only
        self._aborted = threading.Event()
        self._abort_lock = threading.RLock()
        self._abort_reason = ''

        # Set this to a custom Callback() subclass instance before calling
        # strategise() to receive progress reports.
        self._progress_cb = progress.Callback()
        self._tscb = progress.ThreadSafeCallback(self._progress_cb)

        self._exclude_globs = set()  # type: typing.Set[str]

        from blender_asset_tracer.cli import common
        self._shorten = functools.partial(common.shorten, self.project)

        if noop:
            log.warning('Running in no-op mode, only showing what will be done.')

        # Filled by strategise()
        self._actions = collections.defaultdict(AssetAction) \
            # type: typing.DefaultDict[pathlib.Path, AssetAction]
        self.missing_files = set()  # type: typing.Set[pathlib.Path]
        self._new_location_paths = set()  # type: typing.Set[pathlib.Path]
        self._output_path = None  # type: typing.Optional[pathlib.PurePath]

        # Filled by execute()
        self._file_transferer = None  # type: typing.Optional[transfer.FileTransferer]

        # Number of files we would copy, if not for --noop
        self._file_count = 0

        self._tmpdir = tempfile.TemporaryDirectory(prefix='bat-', suffix='-batpack')
        self._rewrite_in = pathlib.Path(self._tmpdir.name)

    def _make_target_path(self, target: str) -> pathlib.PurePath:
        """Return a Path for the given target.

        This can be the target directory itself, but can also be a non-existent
        directory if the target doesn't support direct file access. It should
        only be used to perform path operations, and never for file operations.
        """
        return pathlib.Path(target).absolute()

    def close(self) -> None:
        """Clean up any temporary files."""
        self._tscb.flush()
        self._tmpdir.cleanup()

    def __enter__(self) -> 'Packer':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    @property
    def output_path(self) -> pathlib.PurePath:
        """The path of the packed blend file in the target directory."""
        assert self._output_path is not None
        return self._output_path

    @property
    def progress_cb(self) -> progress.Callback:
        return self._progress_cb

    @progress_cb.setter
    def progress_cb(self, new_progress_cb: progress.Callback):
        self._tscb.flush()
        self._progress_cb = new_progress_cb
        self._tscb = progress.ThreadSafeCallback(self._progress_cb)

    def abort(self, reason='') -> None:
        """Aborts the current packing process.

        Can be called from any thread. Aborts as soon as the running strategise
        or execute function gets control over the execution flow, by raising
        an Aborted exception.
        """
        with self._abort_lock:
            self._abort_reason = reason
            if self._file_transferer:
                self._file_transferer.abort()
            self._aborted.set()

    def _check_aborted(self) -> None:
        """Raises an Aborted exception when abort() was called."""

        with self._abort_lock:
            reason = self._abort_reason
            if self._file_transferer is not None and self._file_transferer.has_error:
                log.error('A transfer error occurred')
                reason = self._file_transferer.error_message()
            elif not self._aborted.is_set():
                return

            log.warning('Aborting')
            self._tscb.flush()
            self._progress_cb.pack_aborted(reason)
            raise Aborted(reason)

    def exclude(self, *globs: str):
        """Register glob-compatible patterns of files that should be ignored.

        Must be called before calling strategise().
        """
        if self._actions:
            raise RuntimeError('%s.exclude() must be called before strategise()' %
                               self.__class__.__qualname__)
        self._exclude_globs.update(globs)

    def strategise(self) -> None:
        """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.

        This function does *not* expand globs. Globs are seen as single
        assets, and are only evaluated when performing the actual transfer
        in the execute() function.
        """

        # The blendfile that we pack is generally not its own dependency, so
        # we have to explicitly add it to the _packed_paths.
        bfile_path = bpathlib.make_absolute(self.blendfile)

        # Both paths have to be resolved first, because this also translates
        # network shares mapped to Windows drive letters back to their UNC
        # notation. Only resolving one but not the other (which can happen
        # with the abosolute() call above) can cause errors.
        bfile_pp = self._target_path / bfile_path.relative_to(bpathlib.make_absolute(self.project))
        self._output_path = bfile_pp

        self._progress_cb.pack_start()

        act = self._actions[bfile_path]
        act.path_action = PathAction.KEEP_PATH
        act.new_path = bfile_pp

        self._check_aborted()
        self._new_location_paths = set()
        for usage in trace.deps(self.blendfile, self._progress_cb):
            self._check_aborted()
            asset_path = usage.abspath
            if any(asset_path.match(glob) for glob in self._exclude_globs):
                log.info('Excluding file: %s', asset_path)
                continue

            if self.relative_only and not usage.asset_path.startswith(b'//'):
                log.info('Skipping absolute path: %s', usage.asset_path)
                continue

            if usage.is_sequence:
                self._visit_sequence(asset_path, usage)
            else:
                self._visit_asset(asset_path, usage)

        self._find_new_paths()
        self._group_rewrites()

    def _visit_sequence(self, asset_path: pathlib.Path, usage: result.BlockUsage):
        assert usage.is_sequence

        for first_path in file_sequence.expand_sequence(asset_path):
            if first_path.exists():
                break
        else:
            # At least the first file of a sequence must exist.
            log.warning('Missing file: %s', asset_path)
            self.missing_files.add(asset_path)
            self._progress_cb.missing_file(asset_path)
            return

        # Handle this sequence as an asset.
        self._visit_asset(asset_path, usage)

    def _visit_asset(self, asset_path: pathlib.Path, usage: result.BlockUsage):
        """Determine what to do with this asset.

        Determines where this asset will be packed, whether it needs rewriting,
        and records the blend file data block referring to it.
        """

        # Sequences are allowed to not exist at this point.
        if not usage.is_sequence and not asset_path.exists():
            log.warning('Missing file: %s', asset_path)
            self.missing_files.add(asset_path)
            self._progress_cb.missing_file(asset_path)
            return

        bfile_path = usage.block.bfile.filepath.absolute()
        self._progress_cb.trace_asset(asset_path)

        # 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).
        if usage.is_sequence:
            first_path = next(file_sequence.expand_sequence(asset_path))
        else:
            first_path = asset_path
        path_in_project = self._path_in_project(first_path)
        use_as_is = usage.asset_path.is_blendfile_relative() and path_in_project
        needs_rewriting = not use_as_is

        act = self._actions[asset_path]
        assert isinstance(act, AssetAction)
        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
            self._new_location_paths.add(asset_path)
        else:
            log.debug('%s can keep using %s', bfile_path, usage.asset_path)
            asset_pp = self._target_path / asset_path.relative_to(self.project)
            act.new_path = asset_pp

    def _find_new_paths(self):
        """Find new locations in the BAT Pack for the given assets."""

        for path in self._new_location_paths:
            act = self._actions[path]
            assert isinstance(act, AssetAction)

            relpath = bpathlib.strip_root(path)
            act.new_path = pathlib.Path(self._target_path, '_outside_project', relpath)

    def _group_rewrites(self) -> None:
        """For each blend file, collect which fields need rewriting.

        This ensures that the execute() step has to visit each blend file
        only once.
        """

        # Take a copy so we can modify self._actions in the loop.
        actions = set(self._actions.values())

        while actions:
            action = actions.pop()

            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 = bpathlib.make_absolute(usage.block.bfile.filepath)
                insert_new_action = bfile_path not in self._actions

                self._actions[bfile_path].rewrites.append(usage)

                if insert_new_action:
                    actions.add(self._actions[bfile_path])

    def _path_in_project(self, path: pathlib.Path) -> bool:
        abs_path = bpathlib.make_absolute(path)
        abs_project = bpathlib.make_absolute(self.project)
        try:
            abs_path.relative_to(abs_project)
        except ValueError:
            return False
        return True

    def execute(self) -> None:
        """Execute the strategy."""
        assert self._actions, 'Run strategise() first'

        if not self.noop:
            self._rewrite_paths()

        self._start_file_transferrer()
        self._perform_file_transfer()
        self._progress_cb.pack_done(self.output_path, self.missing_files)

    def _perform_file_transfer(self):
        """Use file transferrer to do the actual file transfer.

        This is performed in a separate function, so that subclasses can
        override this function to queue up copy/move actions first, and
        then call this function.
        """
        self._write_info_file()
        self._copy_files_to_target()

    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):
        """Starts the file transferrer thread."""
        self._file_transferer = self._create_file_transferer()
        self._file_transferer.progress_cb = self._tscb
        if not self.noop:
            self._file_transferer.start()

    def _copy_files_to_target(self) -> None:
        """Copy all assets to the target directoy.

        This creates the BAT Pack but does not yet do any path rewriting.
        """
        log.debug('Executing %d copy actions', len(self._actions))

        assert self._file_transferer is not None

        try:
            for asset_path, action in self._actions.items():
                self._check_aborted()
                self._copy_asset_and_deps(asset_path, action)

            if self.noop:
                log.info('Would copy %d files to %s', self._file_count, self.target)
                return
            self._file_transferer.done_and_join()
            self._on_file_transfer_finished(file_transfer_completed=True)
        except KeyboardInterrupt:
            log.info('File transfer interrupted with Ctrl+C, aborting.')
            self._file_transferer.abort_and_join()
            self._on_file_transfer_finished(file_transfer_completed=False)
            raise
        finally:
            self._tscb.flush()
            self._check_aborted()

            # Make sure that the file transferer is no longer usable, for
            # example to avoid it being involved in any following call to
            # self.abort().
            self._file_transferer = None

    def _on_file_transfer_finished(self, *, file_transfer_completed: bool) -> None:
        """Called when the file transfer is finished.

        This can be used in subclasses to perform cleanup on the file transferer,
        or to obtain information from it before we destroy it.
        """

    def _rewrite_paths(self) -> None:
        """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
            self._check_aborted()

            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 = action.new_path
            assert bfile_pp is not None

            # 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),
                                                    prefix='bat-',
                                                    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.copy_and_rebind(bfile_tp, mode='rb+')

            for usage in action.rewrites:
                self._check_aborted()
                assert isinstance(usage, result.BlockUsage)
                asset_pp = self._actions[usage.abspath].new_path
                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)
                if usage.path_full_field is None:
                    dir_field = usage.path_dir_field
                    assert dir_field is not None
                    log.debug('   - updating field %s of block %s',
                              dir_field.name.name_only,
                              block)
                    reldir = bpathlib.BlendPath.mkrelative(asset_pp.parent, bfile_pp)
                    written = block.set(dir_field.name.name_only, reldir)
                    log.debug('   - written %d bytes', written)

                    # BIG FAT ASSUMPTION that the filename (e.g. basename
                    # without path) does not change. This makes things much
                    # easier, as in the sequence editor the directory and
                    # filename fields are in different blocks. See the
                    # blocks2assets.scene() function for the implementation.
                else:
                    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.debug('   - written %d bytes', written)

            # Make sure we close the file, otherwise changes may not be
            # flushed before it gets copied.
            if bfile.is_modified:
                self._progress_cb.rewrite_blendfile(bfile_path)
            bfile.close()

    def _copy_asset_and_deps(self, asset_path: pathlib.Path, action: AssetAction):
        # Copy the asset itself, but only if it's not a sequence (sequences are
        # handled below in the for-loop).
        if '*' not in str(asset_path):
            packed_path = action.new_path
            assert packed_path is not None
            read_path = action.read_from or asset_path
            self._send_to_target(read_path, packed_path,
                                 may_move=action.read_from is not None)

        # Copy its sequence dependencies.
        for usage in action.usages:
            if not usage.is_sequence:
                continue

            first_pp = self._actions[usage.abspath].new_path
            assert first_pp is not None

            # In case of globbing, we only support globbing by filename,
            # and not by directory.
            assert '*' not in str(first_pp) or '*' in first_pp.name

            packed_base_dir = first_pp.parent
            for file_path in usage.files():
                packed_path = packed_base_dir / file_path.name
                # Assumption: assets in a sequence are never blend files.
                self._send_to_target(file_path, packed_path)

            # Assumption: all data blocks using this asset use it the same way.
            break

    def _send_to_target(self,
                        asset_path: pathlib.Path,
                        target: pathlib.PurePath,
                        may_move=False):
        if self.noop:
            print('%s -> %s' % (asset_path, target))
            self._file_count += 1
            return

        verb = 'move' if may_move else 'copy'
        log.debug('Queueing %s of %s', verb, asset_path)

        self._tscb.flush()

        assert self._file_transferer is not None
        if may_move:
            self._file_transferer.queue_move(asset_path, target)
        else:
            self._file_transferer.queue_copy(asset_path, target)

    def _write_info_file(self):
        """Write a little text file with info at the top of the pack."""

        infoname = 'pack-info.txt'
        infopath = self._rewrite_in / infoname
        log.debug('Writing info to %s', infopath)
        with infopath.open('wt', encoding='utf8') as infofile:
            print('This is a Blender Asset Tracer pack.', file=infofile)
            print('Start by opening the following blend file:', file=infofile)
            print('    %s' % self._output_path.relative_to(self._target_path).as_posix(),
                  file=infofile)

        self._file_transferer.queue_move(infopath, self._target_path / infoname)