autour_de_minuit/blender_asset_tracer
Archived
8
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial commit

Browse Source
This commit is contained in:
Pullusb
2021-10-18 15:54:04 +02:00
commit 3c5cca81a5
36 changed files with 6461 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pack/__init__.py
+587
View File
@@ -0,0 +1,587 @@
# ***** 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)
pack/filesystem.py
+273
View File
@@ -0,0 +1,273 @@
# ***** 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 logging
import multiprocessing.pool
import pathlib
import shutil
import typing
from .. import compressor
from . import transfer
log = logging.getLogger(__name__)
class AbortTransfer(Exception):
"""Raised when an error was detected and file transfer should be aborted."""
class FileCopier(transfer.FileTransferer):
"""Copies or moves files in source directory order."""
# When we don't compress the files, the process is I/O bound,
# and trashing the storage by using multiple threads will
# only slow things down.
transfer_threads = 1 # type: typing.Optional[int]
def __init__(self):
super().__init__()
self.files_transferred = 0
self.files_skipped = 0
self.already_copied = set()
# (is_dir, action)
self.transfer_funcs = {
(False, transfer.Action.COPY): self.copyfile,
(True, transfer.Action.COPY): self.copytree,
(False, transfer.Action.MOVE): self.move,
(True, transfer.Action.MOVE): self.move,
}
def run(self) -> None:
pool = multiprocessing.pool.ThreadPool(processes=self.transfer_threads)
dst = pathlib.Path()
for src, pure_dst, act in self.iter_queue():
try:
dst = pathlib.Path(pure_dst)
if self.has_error or self._abort.is_set():
raise AbortTransfer()
if self._skip_file(src, dst, act):
continue
# We want to do this in this thread, as it's not thread safe itself.
dst.parent.mkdir(parents=True, exist_ok=True)
pool.apply_async(self._thread, (src, dst, act))
except AbortTransfer:
# either self._error or self._abort is already set. We just have to
# let the system know we didn't handle those files yet.
self.queue.put((src, dst, act), timeout=1.0)
except Exception as ex:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
if self._abort.is_set():
log.debug('Error transferring %s to %s: %s', src, dst, ex)
else:
msg = 'Error transferring %s to %s' % (src, dst)
log.exception(msg)
self.error_set(msg)
# Put the files to copy back into the queue, and abort. This allows
# the main thread to inspect the queue and see which files were not
# copied. The one we just failed (due to this exception) should also
# be reported there.
self.queue.put((src, dst, act), timeout=1.0)
break
log.debug('All transfer threads queued')
pool.close()
log.debug('Waiting for transfer threads to finish')
pool.join()
log.debug('All transfer threads finished')
if self.files_transferred:
log.info('Transferred %d files', self.files_transferred)
if self.files_skipped:
log.info('Skipped %d files', self.files_skipped)
def _thread(self, src: pathlib.Path, dst: pathlib.Path, act: transfer.Action):
try:
tfunc = self.transfer_funcs[src.is_dir(), act]
if self.has_error or self._abort.is_set():
raise AbortTransfer()
log.info('%s %s -> %s', act.name, src, dst)
tfunc(src, dst)
except AbortTransfer:
# either self._error or self._abort is already set. We just have to
# let the system know we didn't handle those files yet.
self.queue.put((src, dst, act), timeout=1.0)
except Exception as ex:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
if self._abort.is_set():
log.debug('Error transferring %s to %s: %s', src, dst, ex)
else:
msg = 'Error transferring %s to %s' % (src, dst)
log.exception(msg)
self.error_set(msg)
# Put the files to copy back into the queue, and abort. This allows
# the main thread to inspect the queue and see which files were not
# copied. The one we just failed (due to this exception) should also
# be reported there.
self.queue.put((src, dst, act), timeout=1.0)
def _skip_file(self, src: pathlib.Path, dst: pathlib.Path, act: transfer.Action) -> bool:
"""Skip this file (return True) or not (return False)."""
st_src = src.stat() # must exist, or it wouldn't be queued.
if not dst.exists():
return False
st_dst = dst.stat()
if st_dst.st_size != st_src.st_size or st_dst.st_mtime < st_src.st_mtime:
return False
log.info('SKIP %s; already exists', src)
if act == transfer.Action.MOVE:
log.debug('Deleting %s', src)
src.unlink()
self.files_skipped += 1
return True
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()
self._move(srcpath, dstpath)
self.files_transferred += 1
self.report_transferred(s_stat.st_size)
def copyfile(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
"""Copy a file, skipping when it already exists."""
if self._abort.is_set() or self.has_error:
return
if (srcpath, dstpath) in self.already_copied:
log.debug('SKIP %s; already copied', srcpath)
return
s_stat = srcpath.stat() # must exist, or it wouldn't be queued.
if dstpath.exists():
d_stat = dstpath.stat()
if d_stat.st_size == s_stat.st_size and d_stat.st_mtime >= s_stat.st_mtime:
log.info('SKIP %s; already exists', srcpath)
self.progress_cb.transfer_file_skipped(srcpath, dstpath)
self.files_skipped += 1
return
log.debug('Copying %s -> %s', srcpath, dstpath)
self._copy(srcpath, dstpath)
self.already_copied.add((srcpath, dstpath))
self.files_transferred += 1
self.report_transferred(s_stat.st_size)
def copytree(self, src: pathlib.Path, dst: pathlib.Path,
symlinks=False, ignore_dangling_symlinks=False):
"""Recursively copy a directory tree.
Copy of shutil.copytree() with some changes:
- Using pathlib
- The destination directory may already exist.
- Existing files with the same file size are skipped.
- Removed ability to ignore things.
"""
if (src, dst) in self.already_copied:
log.debug('SKIP %s; already copied', src)
return
if self.has_error or self._abort.is_set():
raise AbortTransfer()
dst.mkdir(parents=True, exist_ok=True)
errors = [] # type: typing.List[typing.Tuple[pathlib.Path, pathlib.Path, str]]
for srcpath in src.iterdir():
if self.has_error or self._abort.is_set():
raise AbortTransfer()
dstpath = dst / srcpath.name
try:
if srcpath.is_symlink():
linkto = srcpath.resolve()
if symlinks:
# We can't just leave it to `copy_function` because legacy
# code with a custom `copy_function` may rely on copytree
# doing the right thing.
linkto.symlink_to(dstpath)
shutil.copystat(str(srcpath), str(dstpath), follow_symlinks=not symlinks)
else:
# ignore dangling symlink if the flag is on
if not linkto.exists() and ignore_dangling_symlinks:
continue
# otherwise let the copy occurs. copy2 will raise an error
if srcpath.is_dir():
self.copytree(srcpath, dstpath, symlinks)
else:
self.copyfile(srcpath, dstpath)
elif srcpath.is_dir():
self.copytree(srcpath, dstpath, symlinks)
else:
# Will raise a SpecialFileError for unsupported file types
self.copyfile(srcpath, dstpath)
# catch the Error from the recursive copytree so that we can
# continue with other files
except shutil.Error as err:
errors.extend(err.args[0])
except OSError as why:
errors.append((srcpath, dstpath, str(why)))
try:
shutil.copystat(str(src), str(dst))
except OSError as why:
# Copying file access times may fail on Windows
if getattr(why, 'winerror', None) is None:
errors.append((src, dst, str(why)))
if errors:
raise shutil.Error(errors)
self.already_copied.add((src, dst))
return dst
class CompressedFileCopier(FileCopier):
# When we compress the files on the fly, the process is CPU-bound
# so we benefit greatly by multi-threading (packing a Spring scene
# lighting file took 6m30s single-threaded and 2min13 multi-threaded.
transfer_threads = None # type: typing.Optional[int]
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)
pack/progress.py
+148
View File
@@ -0,0 +1,148 @@
# ***** 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
"""Callback class definition for BAT Pack progress reporting."""
import threading
import functools
import logging
import pathlib
import queue
import typing
import blender_asset_tracer.trace.progress
log = logging.getLogger(__name__)
class Callback(blender_asset_tracer.trace.progress.Callback):
"""BAT Pack progress reporting."""
def pack_start(self) -> None:
"""Called when packing starts."""
def pack_done(self,
output_blendfile: pathlib.PurePath,
missing_files: typing.Set[pathlib.Path]) -> None:
"""Called when packing is done."""
def pack_aborted(self, reason: str):
"""Called when packing was aborted."""
def trace_blendfile(self, filename: pathlib.Path) -> None:
"""Called for every blendfile opened when tracing dependencies."""
def trace_asset(self, filename: pathlib.Path) -> None:
"""Called for every asset found when tracing dependencies.
Note that this can also be a blend file.
"""
def rewrite_blendfile(self, orig_filename: pathlib.Path) -> None:
"""Called for every rewritten blendfile."""
def transfer_file(self, src: pathlib.Path, dst: pathlib.PurePath) -> None:
"""Called when a file transfer starts."""
def transfer_file_skipped(self, src: pathlib.Path, dst: pathlib.PurePath) -> None:
"""Called when a file is skipped because it already exists."""
def transfer_progress(self, total_bytes: int, transferred_bytes: int) -> None:
"""Called during file transfer, with per-pack info (not per file).
:param total_bytes: The total amount of bytes to be transferred for
the current packing operation. This can increase while transfer
is happening, when more files are discovered (because transfer
starts in a separate thread before all files are found).
:param transferred_bytes: The total amount of bytes transfered for
the current packing operation.
"""
def missing_file(self, filename: pathlib.Path) -> None:
"""Called for every asset that does not exist on the filesystem."""
class ThreadSafeCallback(Callback):
"""Thread-safe wrapper for Callback instances.
Progress calls are queued until flush() is called. The queued calls are
called in the same thread as the one calling flush().
"""
def __init__(self, wrapped: Callback) -> None:
self.log = log.getChild('ThreadSafeCallback')
self.wrapped = wrapped
# Thread-safe queue for passing progress reports on the main thread.
self._reporting_queue = queue.Queue() # type: queue.Queue[typing.Callable]
self._main_thread_id = threading.get_ident()
def _queue(self, func: typing.Callable, *args, **kwargs):
partial = functools.partial(func, *args, **kwargs)
if self._main_thread_id == threading.get_ident():
partial()
else:
self._reporting_queue.put(partial)
def pack_start(self) -> None:
self._queue(self.wrapped.pack_start)
def pack_done(self,
output_blendfile: pathlib.PurePath,
missing_files: typing.Set[pathlib.Path]) -> None:
self._queue(self.wrapped.pack_done, output_blendfile, missing_files)
def pack_aborted(self, reason: str):
self._queue(self.wrapped.pack_aborted, reason)
def trace_blendfile(self, filename: pathlib.Path) -> None:
self._queue(self.wrapped.trace_blendfile, filename)
def trace_asset(self, filename: pathlib.Path) -> None:
self._queue(self.wrapped.trace_asset, filename)
def transfer_file(self, src: pathlib.Path, dst: pathlib.PurePath) -> None:
self._queue(self.wrapped.transfer_file, src, dst)
def transfer_file_skipped(self, src: pathlib.Path, dst: pathlib.PurePath) -> None:
self._queue(self.wrapped.transfer_file_skipped, src, dst)
def transfer_progress(self, total_bytes: int, transferred_bytes: int) -> None:
self._queue(self.wrapped.transfer_progress, total_bytes, transferred_bytes)
def missing_file(self, filename: pathlib.Path) -> None:
self._queue(self.wrapped.missing_file, filename)
def flush(self, timeout: float = None) -> None:
"""Call the queued calls, call this in the main thread."""
while True:
try:
call = self._reporting_queue.get(block=timeout is not None,
timeout=timeout)
except queue.Empty:
return
try:
call()
except Exception:
# Don't let the handling of one callback call
# block the entire flush process.
self.log.exception('Error calling %s', call)
pack/s3.py
+182
View File
@@ -0,0 +1,182 @@
# ***** 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
"""Amazon S3-compatible uploader."""
import hashlib
import logging
import pathlib
import typing
import urllib.parse
from . import Packer, transfer
log = logging.getLogger(__name__)
# TODO(Sybren): compute MD5 sums of queued files in a separate thread, so that
# we can upload a file to S3 and compute an MD5 of another file simultaneously.
def compute_md5(filepath: pathlib.Path) -> str:
log.debug('Computing MD5sum of %s', filepath)
hasher = hashlib.md5()
with filepath.open('rb') as infile:
while True:
block = infile.read(102400)
if not block:
break
hasher.update(block)
md5 = hasher.hexdigest()
log.debug('MD5sum of %s is %s', filepath, md5)
return md5
class S3Packer(Packer):
"""Creates BAT Packs on S3-compatible storage."""
def __init__(self, *args, endpoint, **kwargs) -> None:
"""Constructor
:param endpoint: URL of the S3 storage endpoint
"""
super().__init__(*args, **kwargs)
import boto3
# Create a session so that credentials can be read from the [endpoint]
# section in ~/.aws/credentials.
# See https://boto3.readthedocs.io/en/latest/guide/configuration.html#guide-configuration
components = urllib.parse.urlparse(endpoint)
profile_name = components.netloc
endpoint = urllib.parse.urlunparse(components)
log.debug('Using Boto3 profile name %r for url %r', profile_name, endpoint)
self.session = boto3.Session(profile_name=profile_name)
self.client = self.session.client('s3', endpoint_url=endpoint)
def set_credentials(self,
endpoint: str,
access_key_id: str,
secret_access_key: str):
"""Set S3 credentials."""
self.client = self.session.client('s3',
endpoint_url=endpoint,
aws_access_key_id=access_key_id,
aws_secret_access_key=secret_access_key)
def _create_file_transferer(self) -> transfer.FileTransferer:
return S3Transferrer(self.client)
class S3Transferrer(transfer.FileTransferer):
"""Copies or moves files in source directory order."""
class AbortUpload(Exception):
"""Raised from the upload callback to abort an upload."""
def __init__(self, botoclient) -> None:
super().__init__()
self.client = botoclient
def run(self) -> None:
files_transferred = 0
files_skipped = 0
for src, dst, act in self.iter_queue():
try:
did_upload = self.upload_file(src, dst)
files_transferred += did_upload
files_skipped += not did_upload
if act == transfer.Action.MOVE:
self.delete_file(src)
except Exception:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
log.exception('Error transferring %s to %s', src, dst)
# Put the files to copy back into the queue, and abort. This allows
# the main thread to inspect the queue and see which files were not
# copied. The one we just failed (due to this exception) should also
# be reported there.
self.queue.put((src, dst, act))
return
if files_transferred:
log.info('Transferred %d files', files_transferred)
if files_skipped:
log.info('Skipped %d files', files_skipped)
def upload_file(self, src: pathlib.Path, dst: pathlib.PurePath) -> bool:
"""Upload a file to an S3 bucket.
The first part of 'dst' is used as the bucket name, the remained as the
path inside the bucket.
:returns: True if the file was uploaded, False if it was skipped.
"""
bucket = dst.parts[0]
dst_path = pathlib.Path(*dst.parts[1:])
md5 = compute_md5(src)
key = str(dst_path)
existing_md5, existing_size = self.get_metadata(bucket, key)
if md5 == existing_md5 and src.stat().st_size == existing_size:
log.debug('skipping %s, it already exists on the server with MD5 %s',
src, existing_md5)
return False
log.info('Uploading %s', src)
try:
self.client.upload_file(str(src),
Bucket=bucket,
Key=key,
Callback=self.report_transferred,
ExtraArgs={'Metadata': {'md5': md5}})
except self.AbortUpload:
return False
return True
def report_transferred(self, bytes_transferred: int):
if self._abort.is_set():
log.warning('Interrupting ongoing upload')
raise self.AbortUpload('interrupting ongoing upload')
super().report_transferred(bytes_transferred)
def get_metadata(self, bucket: str, key: str) -> typing.Tuple[str, int]:
"""Get MD5 sum and size on S3.
:returns: the MD5 hexadecimal hash and the file size in bytes.
If the file does not exist or has no known MD5 sum,
returns ('', -1)
"""
import botocore.exceptions
log.debug('Getting metadata of %s/%s', bucket, key)
try:
info = self.client.head_object(Bucket=bucket, Key=key)
except botocore.exceptions.ClientError as ex:
error_code = ex.response.get('Error').get('Code', 'Unknown')
# error_code already is a string, but this makes the code forward
# compatible with a time where they use integer codes.
if str(error_code) == '404':
return '', -1
raise ValueError('error response:' % ex.response) from None
try:
return info['Metadata']['md5'], info['ContentLength']
except KeyError:
return '', -1
pack/shaman/__init__.py
+130
View File
@@ -0,0 +1,130 @@
# ***** 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) 2019, Blender Foundation - Sybren A. Stüvel
"""Shaman Client interface."""
import logging
import os
import pathlib
import typing
import urllib.parse
import requests
import blender_asset_tracer.pack as bat_pack
import blender_asset_tracer.pack.transfer as bat_transfer
from .transfer import ShamanTransferrer
from .client import ShamanClient
log = logging.getLogger(__name__)
class ShamanPacker(bat_pack.Packer):
"""Creates BAT Packs on a Shaman server."""
def __init__(self,
bfile: pathlib.Path,
project: pathlib.Path,
target: str,
endpoint: str,
checkout_id: str,
**kwargs) -> None:
"""Constructor
:param target: mock target '/' to construct project-relative paths.
:param endpoint: URL of the Shaman endpoint.
"""
super().__init__(bfile, project, target, **kwargs)
self.checkout_id = checkout_id
self.shaman_endpoint = endpoint
self._checkout_location = ''
def _get_auth_token(self) -> str:
# TODO: get a token from the Flamenco Server.
token_from_env = os.environ.get('SHAMAN_JWT_TOKEN')
if token_from_env:
return token_from_env
log.warning('Using temporary hack to get auth token from Shaman, '
'set SHAMAN_JTW_TOKEN to prevent')
unauth_shaman = ShamanClient('', self.shaman_endpoint)
resp = unauth_shaman.get('get-token', timeout=10)
resp.raise_for_status()
return resp.text
def _create_file_transferer(self) -> bat_transfer.FileTransferer:
# TODO: pass self._get_auth_token itself, so that the Transferer will be able to
# decide when to get this token (and how many times).
auth_token = self._get_auth_token()
return ShamanTransferrer(auth_token, self.project, self.shaman_endpoint, self.checkout_id)
def _make_target_path(self, target: str) -> pathlib.PurePath:
return pathlib.PurePosixPath('/')
def _on_file_transfer_finished(self, *, file_transfer_completed: bool):
super()._on_file_transfer_finished(file_transfer_completed=file_transfer_completed)
assert isinstance(self._file_transferer, ShamanTransferrer)
self._checkout_location = self._file_transferer.checkout_location
@property
def checkout_location(self) -> str:
"""Return the checkout location of the packed blend file.
:return: the checkout location, or '' if no checkout was made.
"""
return self._checkout_location
@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
checkout_location = pathlib.PurePosixPath(self._checkout_location)
rel_output = self._output_path.relative_to(self._target_path)
return checkout_location / rel_output
def execute(self):
try:
super().execute()
except requests.exceptions.ConnectionError as ex:
log.exception('Error communicating with Shaman')
self.abort(str(ex))
self._check_aborted()
def parse_endpoint(shaman_url: str) -> typing.Tuple[str, str]:
"""Convert shaman://hostname/path#checkoutID into endpoint URL + checkout ID."""
urlparts = urllib.parse.urlparse(str(shaman_url))
if urlparts.scheme in {'shaman', 'shaman+https'}:
scheme = 'https'
elif urlparts.scheme == 'shaman+http':
scheme = 'http'
else:
raise ValueError('Invalid scheme %r, choose shaman:// or shaman+http://', urlparts.scheme)
checkout_id = urllib.parse.unquote(urlparts.fragment)
path = urlparts.path or '/'
new_urlparts = (scheme, urlparts.netloc, path, *urlparts[3:-1], '')
endpoint = urllib.parse.urlunparse(new_urlparts)
return endpoint, checkout_id
pack/shaman/cache.py
+197
View File
@@ -0,0 +1,197 @@
# ***** 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) 2019, Blender Foundation - Sybren A. Stüvel
import base64
import hashlib
import json
import logging
import sys
import time
import typing
from collections import deque
from pathlib import Path
from . import time_tracker
CACHE_ROOT = Path().home() / '.cache/shaman-client/shasums'
MAX_CACHE_FILES_AGE_SECS = 3600 * 24 * 60 # 60 days
log = logging.getLogger(__name__)
class TimeInfo:
computing_checksums = 0.0
checksum_cache_handling = 0.0
def find_files(root: Path) -> typing.Iterable[Path]:
"""Recursively finds files in the given root path.
Directories are recursed into, and file paths are yielded.
Symlinks are yielded if they refer to a regular file.
"""
queue = deque([root])
while queue:
path = queue.popleft()
# Ignore hidden files/dirs; these can be things like '.svn' or '.git',
# which shouldn't be sent to Shaman.
if path.name.startswith('.'):
continue
if path.is_dir():
for child in path.iterdir():
queue.append(child)
continue
# Only yield symlinks if they link to (a link to) a normal file.
if path.is_symlink():
symlinked = path.resolve()
if symlinked.is_file():
yield path
continue
if path.is_file():
yield path
def compute_checksum(filepath: Path) -> str:
"""Compute the SHA256 checksum for the given file."""
blocksize = 32 * 1024
log.debug('Computing checksum of %s', filepath)
with time_tracker.track_time(TimeInfo, 'computing_checksums'):
hasher = hashlib.sha256()
with filepath.open('rb') as infile:
while True:
block = infile.read(blocksize)
if not block:
break
hasher.update(block)
checksum = hasher.hexdigest()
return checksum
def _cache_path(filepath: Path) -> Path:
"""Compute the cache file for the given file path."""
fs_encoding = sys.getfilesystemencoding()
filepath = filepath.absolute()
# Reverse the directory, because most variation is in the last bytes.
rev_dir = str(filepath.parent)[::-1]
encoded_path = filepath.stem + rev_dir + filepath.suffix
cache_key = base64.urlsafe_b64encode(encoded_path.encode(fs_encoding)).decode().rstrip('=')
cache_path = CACHE_ROOT / cache_key[:10] / cache_key[10:]
return cache_path
def compute_cached_checksum(filepath: Path) -> str:
"""Computes the SHA256 checksum.
The checksum is cached to disk. If the cache is still valid, it is used to
skip the actual SHA256 computation.
"""
with time_tracker.track_time(TimeInfo, 'checksum_cache_handling'):
current_stat = filepath.stat()
cache_path = _cache_path(filepath)
try:
with cache_path.open('r') as cache_file:
payload = json.load(cache_file)
except (OSError, ValueError):
# File may not exist, or have invalid contents.
pass
else:
checksum = payload.get('checksum', '')
cached_mtime = payload.get('file_mtime', 0.0)
cached_size = payload.get('file_size', -1)
if (checksum
and current_stat.st_size == cached_size
and abs(cached_mtime - current_stat.st_mtime) < 0.01):
cache_path.touch()
return checksum
checksum = compute_checksum(filepath)
with time_tracker.track_time(TimeInfo, 'checksum_cache_handling'):
payload = {
'checksum': checksum,
'file_mtime': current_stat.st_mtime,
'file_size': current_stat.st_size,
}
try:
cache_path.parent.mkdir(parents=True, exist_ok=True)
with cache_path.open('w') as cache_file:
json.dump(payload, cache_file)
except IOError as ex:
log.warning('Unable to write checksum cache file %s: %s', cache_path, ex)
return checksum
def cleanup_cache() -> None:
"""Remove all cache files that are older than MAX_CACHE_FILES_AGE_SECS."""
if not CACHE_ROOT.exists():
return
with time_tracker.track_time(TimeInfo, 'checksum_cache_handling'):
queue = deque([CACHE_ROOT])
rmdir_queue = []
now = time.time()
num_removed_files = 0
num_removed_dirs = 0
while queue:
path = queue.popleft()
if path.is_dir():
queue.extend(path.iterdir())
rmdir_queue.append(path)
continue
assert path.is_file()
path.relative_to(CACHE_ROOT)
age = now - path.stat().st_mtime
# Don't trust files from the future either.
if 0 <= age <= MAX_CACHE_FILES_AGE_SECS:
continue
path.unlink()
num_removed_files += 1
for dirpath in reversed(rmdir_queue):
assert dirpath.is_dir()
dirpath.relative_to(CACHE_ROOT)
try:
dirpath.rmdir()
num_removed_dirs += 1
except OSError:
pass
if num_removed_dirs or num_removed_files:
log.info('Cache Cleanup: removed %d dirs and %d files', num_removed_dirs, num_removed_files)
pack/shaman/client.py
+129
View File
@@ -0,0 +1,129 @@
# ***** 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) 2019, Blender Foundation - Sybren A. Stüvel
import urllib.parse
import requests.packages.urllib3.util.retry
import requests.adapters
class ShamanClient:
"""Thin wrapper around a Requests session to perform Shaman requests."""
def __init__(self, auth_token: str, base_url: str):
self._auth_token = auth_token
self._base_url = base_url
retries = requests.packages.urllib3.util.retry.Retry(
total=10,
backoff_factor=0.05,
)
http_adapter = requests.adapters.HTTPAdapter(max_retries=retries)
self._session = requests.session()
self._session.mount('https://', http_adapter)
self._session.mount('http://', http_adapter)
if auth_token:
self._session.headers['Authorization'] = 'Bearer ' + auth_token
def request(self, method: str, url: str, **kwargs) -> requests.Response:
kwargs.setdefault('timeout', 300)
full_url = urllib.parse.urljoin(self._base_url, url)
return self._session.request(method, full_url, **kwargs)
def get(self, url, **kwargs):
r"""Sends a GET request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', True)
return self.request('GET', url, **kwargs)
def options(self, url, **kwargs):
r"""Sends a OPTIONS request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', True)
return self.request('OPTIONS', url, **kwargs)
def head(self, url, **kwargs):
r"""Sends a HEAD request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', False)
return self.request('HEAD', url, **kwargs)
def post(self, url, data=None, json=None, **kwargs):
r"""Sends a POST request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param data: (optional) Dictionary, list of tuples, bytes, or file-like
object to send in the body of the :class:`Request`.
:param json: (optional) json to send in the body of the :class:`Request`.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
return self.request('POST', url, data=data, json=json, **kwargs)
def put(self, url, data=None, **kwargs):
r"""Sends a PUT request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param data: (optional) Dictionary, list of tuples, bytes, or file-like
object to send in the body of the :class:`Request`.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
return self.request('PUT', url, data=data, **kwargs)
def patch(self, url, data=None, **kwargs):
r"""Sends a PATCH request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param data: (optional) Dictionary, list of tuples, bytes, or file-like
object to send in the body of the :class:`Request`.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
return self.request('PATCH', url, data=data, **kwargs)
def delete(self, url, **kwargs):
r"""Sends a DELETE request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
return self.request('DELETE', url, **kwargs)
pack/shaman/time_tracker.py
+32
View File
@@ -0,0 +1,32 @@
# ***** 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) 2019, Blender Foundation - Sybren A. Stüvel
import contextlib
import time
import typing
@contextlib.contextmanager
def track_time(tracker_object: typing.Any, attribute: str):
"""Context manager, tracks how long the context took to run."""
start_time = time.monotonic()
yield
duration = time.monotonic() - start_time
tracked_so_far = getattr(tracker_object, attribute, 0.0)
setattr(tracker_object, attribute, tracked_so_far + duration)
pack/shaman/transfer.py
+359
View File
@@ -0,0 +1,359 @@
# ***** 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) 2019, Blender Foundation - Sybren A. Stüvel
import collections
import logging
import pathlib
import random
import typing
import requests
import blender_asset_tracer.pack.transfer as bat_transfer
from blender_asset_tracer import bpathlib
MAX_DEFERRED_PATHS = 8
MAX_FAILED_PATHS = 8
response_file_unknown = "file-unknown"
response_already_uploading = "already-uploading"
class FileInfo:
def __init__(self, checksum: str, filesize: int, abspath: pathlib.Path):
self.checksum = checksum
self.filesize = filesize
self.abspath = abspath
class ShamanTransferrer(bat_transfer.FileTransferer):
"""Sends files to a Shaman server."""
class AbortUpload(Exception):
"""Raised from the upload callback to abort an upload."""
def __init__(self, auth_token: str, project_root: pathlib.Path,
shaman_endpoint: str, checkout_id: str) -> None:
from . import client
super().__init__()
self.client = client.ShamanClient(auth_token, shaman_endpoint)
self.project_root = project_root
self.checkout_id = checkout_id
self.log = logging.getLogger(__name__)
self._file_info = {} # type: typing.Dict[str, FileInfo]
# When the Shaman creates a checkout, it'll return the location of that
# checkout. This can then be combined with the project-relative path
# of the to-be-rendered blend file (e.g. the one 'bat pack' was pointed
# at).
self._checkout_location = ''
self.uploaded_files = 0
self.uploaded_bytes = 0
# noinspection PyBroadException
def run(self) -> None:
try:
self.uploaded_files = 0
self.uploaded_bytes = 0
# Construct the Shaman Checkout Definition file.
# This blocks until we know the entire list of files to transfer.
definition_file, allowed_relpaths, delete_when_done = self._create_checkout_definition()
if not definition_file:
# An error has already been logged.
return
self.log.info('Created checkout definition file of %d KiB',
len(definition_file) // 1024)
self.log.info('Feeding %d files to the Shaman', len(self._file_info))
if self.log.isEnabledFor(logging.INFO):
for path in self._file_info:
self.log.info(' - %s', path)
# Try to upload all the files.
failed_paths = set() # type: typing.Set[str]
max_tries = 50
for try_index in range(max_tries):
# Send the file to the Shaman and see what we still need to send there.
to_upload = self._send_checkout_def_to_shaman(definition_file, allowed_relpaths)
if to_upload is None:
# An error has already been logged.
return
if not to_upload:
break
# Send the files that still need to be sent.
self.log.info('Upload attempt %d', try_index + 1)
failed_paths = self._upload_files(to_upload)
if not failed_paths:
break
# Having failed paths at this point is expected when multiple
# clients are sending the same files. Instead of retrying on a
# file-by-file basis, we just re-send the checkout definition
# file to the Shaman and obtain a new list of files to upload.
if failed_paths:
self.log.error('Aborting upload due to too many failures')
self.error_set('Giving up after %d attempts to upload the files' % max_tries)
return
self.log.info('All files uploaded succesfully')
self._request_checkout(definition_file)
# Delete the files that were supposed to be moved.
for src in delete_when_done:
self.delete_file(src)
except Exception as ex:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
self.log.exception('Error transferring files to Shaman')
self.error_set('Unexpected exception transferring files to Shaman: %s' % ex)
# noinspection PyBroadException
def _create_checkout_definition(self) \
-> typing.Tuple[bytes, typing.Set[str], typing.List[pathlib.Path]]:
"""Create the checkout definition file for this BAT pack.
:returns: the checkout definition (as bytes), a set of paths in that file,
and list of paths to delete.
If there was an error and file transfer was aborted, the checkout
definition file will be empty.
"""
from . import cache
definition_lines = [] # type: typing.List[bytes]
delete_when_done = [] # type: typing.List[pathlib.Path]
# We keep track of the relative paths we want to send to the Shaman,
# so that the Shaman cannot ask us to upload files we didn't want to.
relpaths = set() # type: typing.Set[str]
for src, dst, act in self.iter_queue():
try:
checksum = cache.compute_cached_checksum(src)
filesize = src.stat().st_size
# relpath = dst.relative_to(self.project_root)
relpath = bpathlib.strip_root(dst).as_posix()
self._file_info[relpath] = FileInfo(
checksum=checksum,
filesize=filesize,
abspath=src,
)
line = '%s %s %s' % (checksum, filesize, relpath)
definition_lines.append(line.encode('utf8'))
relpaths.add(relpath)
if act == bat_transfer.Action.MOVE:
delete_when_done.append(src)
except Exception:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
msg = 'Error transferring %s to %s' % (src, dst)
self.log.exception(msg)
# Put the files to copy back into the queue, and abort. This allows
# the main thread to inspect the queue and see which files were not
# copied. The one we just failed (due to this exception) should also
# be reported there.
self.queue.put((src, dst, act))
self.error_set(msg)
return b'', set(), delete_when_done
cache.cleanup_cache()
return b'\n'.join(definition_lines), relpaths, delete_when_done
def _send_checkout_def_to_shaman(self, definition_file: bytes,
allowed_relpaths: typing.Set[str]) \
-> typing.Optional[collections.deque]:
"""Send the checkout definition file to the Shaman.
:return: An iterable of paths (relative to the project root) that still
need to be uploaded, or None if there was an error.
"""
resp = self.client.post('checkout/requirements', data=definition_file, stream=True,
headers={'Content-Type': 'text/plain'},
timeout=15)
if resp.status_code >= 300:
msg = 'Error from Shaman, code %d: %s' % (resp.status_code, resp.text)
self.log.error(msg)
self.error_set(msg)
return None
to_upload = collections.deque() # type: collections.deque
for line in resp.iter_lines():
response, path = line.decode().split(' ', 1)
self.log.debug(' %s: %s', response, path)
if path not in allowed_relpaths:
msg = 'Shaman requested path we did not intend to upload: %r' % path
self.log.error(msg)
self.error_set(msg)
return None
if response == response_file_unknown:
to_upload.appendleft(path)
elif response == response_already_uploading:
to_upload.append(path)
elif response == 'ERROR':
msg = 'Error from Shaman: %s' % path
self.log.error(msg)
self.error_set(msg)
return None
else:
msg = 'Unknown response from Shaman for path %r: %r' % (path, response)
self.log.error(msg)
self.error_set(msg)
return None
return to_upload
def _upload_files(self, to_upload: collections.deque) -> typing.Set[str]:
"""Actually upload the files to Shaman.
Returns the set of files that we did not upload.
"""
failed_paths = set() # type: typing.Set[str]
deferred_paths = set()
def defer(some_path: str):
nonlocal to_upload
self.log.info(' %s deferred (already being uploaded by someone else)', some_path)
deferred_paths.add(some_path)
# Instead of deferring this one file, randomize the files to upload.
# This prevents multiple deferrals when someone else is uploading
# files from the same project (because it probably happens alphabetically).
all_files = list(to_upload)
random.shuffle(all_files)
to_upload = collections.deque(all_files)
if not to_upload:
self.log.info('All %d files are at the Shaman already', len(self._file_info))
self.report_transferred(0)
return failed_paths
self.log.info('Going to upload %d of %d files', len(to_upload), len(self._file_info))
while to_upload:
# After too many failures, just retry to get a fresh set of files to upload.
if len(failed_paths) > MAX_FAILED_PATHS:
self.log.info('Too many failures, going to abort this iteration')
failed_paths.update(to_upload)
return failed_paths
path = to_upload.popleft()
fileinfo = self._file_info[path]
self.log.info(' %s', path)
headers = {
'X-Shaman-Original-Filename': path,
}
# Let the Shaman know whether we can defer uploading this file or not.
can_defer = (len(deferred_paths) < MAX_DEFERRED_PATHS
and path not in deferred_paths
and len(to_upload))
if can_defer:
headers['X-Shaman-Can-Defer-Upload'] = 'true'
url = 'files/%s/%d' % (fileinfo.checksum, fileinfo.filesize)
try:
with fileinfo.abspath.open('rb') as infile:
resp = self.client.post(url, data=infile, headers=headers)
except requests.ConnectionError as ex:
if can_defer:
# Closing the connection with an 'X-Shaman-Can-Defer-Upload: true' header
# indicates that we should defer the upload. Requests doesn't give us the
# reply, even though it was written by the Shaman before it closed the
# connection.
defer(path)
else:
self.log.info(' %s could not be uploaded, might retry later: %s', path, ex)
failed_paths.add(path)
continue
if resp.status_code == 208:
# For small files we get the 208 response, because the server closes the
# connection after we sent the entire request. For bigger files the server
# responds sooner, and Requests gives us the above ConnectionError.
if can_defer:
defer(path)
else:
self.log.info(' %s skipped (already existed on the server)', path)
continue
if resp.status_code >= 300:
msg = 'Error from Shaman uploading %s, code %d: %s' % (
fileinfo.abspath, resp.status_code, resp.text)
self.log.error(msg)
self.error_set(msg)
return failed_paths
failed_paths.discard(path)
self.uploaded_files += 1
file_size = fileinfo.abspath.stat().st_size
self.uploaded_bytes += file_size
self.report_transferred(file_size)
if not failed_paths:
self.log.info('Done uploading %d bytes in %d files',
self.uploaded_bytes, self.uploaded_files)
else:
self.log.info('Uploaded %d bytes in %d files so far',
self.uploaded_bytes, self.uploaded_files)
return failed_paths
def report_transferred(self, bytes_transferred: int):
if self._abort.is_set():
self.log.warning('Interrupting ongoing upload')
raise self.AbortUpload('interrupting ongoing upload')
super().report_transferred(bytes_transferred)
def _request_checkout(self, definition_file: bytes):
"""Ask the Shaman to create a checkout of this BAT pack."""
if not self.checkout_id:
self.log.warning('NOT requesting checkout at Shaman')
return
self.log.info('Requesting checkout at Shaman for checkout_id=%r', self.checkout_id)
resp = self.client.post('checkout/create/%s' % self.checkout_id, data=definition_file,
headers={'Content-Type': 'text/plain'})
if resp.status_code >= 300:
msg = 'Error from Shaman, code %d: %s' % (resp.status_code, resp.text)
self.log.error(msg)
self.error_set(msg)
return
self._checkout_location = resp.text.strip()
self.log.info('Response from Shaman, code %d: %s', resp.status_code, resp.text)
@property
def checkout_location(self) -> str:
"""Returns the checkout location, or '' if no checkout was made."""
if not self._checkout_location:
return ''
return self._checkout_location
pack/transfer.py
+221
View File
@@ -0,0 +1,221 @@
# ***** 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 abc
import enum
import logging
import pathlib
import queue
import threading
import time
import typing
from . import progress
log = logging.getLogger(__name__)
class FileTransferError(IOError):
"""Raised when one or more files could not be transferred."""
def __init__(self, message, files_remaining: typing.List[pathlib.Path]) -> None:
super().__init__(message)
self.files_remaining = files_remaining
class Action(enum.Enum):
COPY = 1
MOVE = 2
QueueItem = typing.Tuple[pathlib.Path, pathlib.PurePath, Action]
class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
"""Abstract superclass for file transfer classes.
Implement a run() function in a subclass that performs the actual file
transfer.
"""
def __init__(self) -> None:
super().__init__()
self.log = log.getChild('FileTransferer')
# For copying in a different process. By using a priority queue the files
# are automatically sorted alphabetically, which means we go through all files
# in a single directory at a time. This should be faster to copy than random
# access. The order isn't guaranteed, though, as we're not waiting around for
# all file paths to be known before copying starts.
# maxsize=100 is just a guess as to a reasonable upper limit. When this limit
# is reached, the main thread will simply block while waiting for this thread
# to finish copying a file.
self.queue = queue.PriorityQueue(maxsize=100) # type: queue.PriorityQueue[QueueItem]
self.done = threading.Event()
self._abort = threading.Event() # Indicates user-requested abort
self.__error_mutex = threading.Lock()
self.__error = threading.Event() # Indicates abort due to some error
self.__error_message = ''
# Instantiate a dummy progress callback so that we can call it
# without checking for None all the time.
self.progress_cb = progress.ThreadSafeCallback(progress.Callback())
self.total_queued_bytes = 0
self.total_transferred_bytes = 0
@abc.abstractmethod
def run(self):
"""Perform actual file transfer in a thread."""
def queue_copy(self, src: pathlib.Path, dst: pathlib.PurePath):
"""Queue a copy action from 'src' to 'dst'."""
assert not self.done.is_set(), 'Queueing not allowed after done_and_join() was called'
assert not self._abort.is_set(), 'Queueing not allowed after abort_and_join() was called'
if self.__error.is_set():
return
self.queue.put((src, dst, Action.COPY))
self.total_queued_bytes += src.stat().st_size
def queue_move(self, src: pathlib.Path, dst: pathlib.PurePath):
"""Queue a move action from 'src' to 'dst'."""
assert not self.done.is_set(), 'Queueing not allowed after done_and_join() was called'
assert not self._abort.is_set(), 'Queueing not allowed after abort_and_join() was called'
if self.__error.is_set():
return
self.queue.put((src, dst, Action.MOVE))
self.total_queued_bytes += src.stat().st_size
def report_transferred(self, bytes_transferred: int):
"""Report transfer of `block_size` bytes."""
self.total_transferred_bytes += bytes_transferred
self.progress_cb.transfer_progress(self.total_queued_bytes, self.total_transferred_bytes)
def done_and_join(self) -> None:
"""Indicate all files have been queued, and wait until done.
After this function has been called, the queue_xxx() methods should not
be called any more.
:raises FileTransferError: if there was an error transferring one or
more files.
"""
self.done.set()
self.join()
if not self.queue.empty():
# Flush the queue so that we can report which files weren't copied yet.
files_remaining = self._files_remaining()
assert files_remaining
raise FileTransferError(
"%d files couldn't be transferred" % len(files_remaining),
files_remaining)
def _files_remaining(self) -> typing.List[pathlib.Path]:
"""Source files that were queued but not transferred."""
files_remaining = []
while not self.queue.empty():
src, dst, act = self.queue.get_nowait()
files_remaining.append(src)
return files_remaining
def abort(self) -> None:
"""Abort the file transfer, immediately returns."""
log.info('Aborting')
self._abort.set()
def abort_and_join(self) -> None:
"""Abort the file transfer, and wait until done."""
self.abort()
self.join()
files_remaining = self._files_remaining()
if not files_remaining:
return
log.warning("%d files couldn't be transferred, starting with %s",
len(files_remaining), files_remaining[0])
def iter_queue(self) -> typing.Iterable[QueueItem]:
"""Generator, yield queued items until the work is done."""
while True:
if self._abort.is_set() or self.__error.is_set():
return
try:
src, dst, action = self.queue.get(timeout=0.5)
self.progress_cb.transfer_file(src, dst)
yield src, dst, action
except queue.Empty:
if self.done.is_set():
return
def join(self, timeout: float = None) -> None:
"""Wait for the transfer to finish/stop."""
if timeout:
run_until = time.time() + timeout
else:
run_until = float('inf')
# We can't simply block the thread, we have to keep watching the
# progress queue.
while self.is_alive():
if time.time() > run_until:
self.log.warning('Timeout while waiting for transfer to finish')
return
self.progress_cb.flush(timeout=0.5)
# Since Thread.join() neither returns anything nor raises any exception
# when timing out, we don't even have to call it any more.
def delete_file(self, path: pathlib.Path):
"""Deletes a file, only logging a warning if deletion fails."""
log.debug('Deleting %s, file has been transferred', path)
try:
path.unlink()
except IOError as ex:
log.warning('Unable to delete %s: %s', path, ex)
@property
def has_error(self) -> bool:
return self.__error.is_set()
def error_set(self, message: str):
"""Indicate an error occurred, and provide a message."""
with self.__error_mutex:
# Avoid overwriting previous error messages.
if self.__error.is_set():
return
self.__error.set()
self.__error_message = message
def error_message(self) -> str:
"""Retrieve the error messsage, or an empty string if no error occurred."""
with self.__error_mutex:
if not self.__error.is_set():
return ''
return self.__error_message
pack/zipped.py
+89
View File
@@ -0,0 +1,89 @@
# ***** 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
"""ZIP file packer.
Note: There is no official file name encoding for ZIP files. Expect trouble
when you want to use the ZIP cross-platform and you have non-ASCII names.
"""
import logging
import pathlib
from . import Packer, transfer
log = logging.getLogger(__name__)
# Suffixes to store uncompressed in the zip.
STORE_ONLY = {'.jpg', '.jpeg', '.exr'}
class ZipPacker(Packer):
"""Creates a zipped BAT Pack instead of a directory."""
def _create_file_transferer(self) -> transfer.FileTransferer:
target_path = pathlib.Path(self._target_path)
return ZipTransferrer(target_path.absolute())
class ZipTransferrer(transfer.FileTransferer):
"""Creates a ZIP file instead of writing to a directory.
Note: There is no official file name encoding for ZIP files. If you have
unicode file names, they will be encoded as UTF-8. WinZip interprets all
file names as encoded in CP437, also known as DOS Latin.
"""
def __init__(self, zippath: pathlib.Path) -> None:
super().__init__()
self.zippath = zippath
def run(self) -> None:
import zipfile
zippath = self.zippath.absolute()
with zipfile.ZipFile(str(zippath), 'w') as outzip:
for src, dst, act in self.iter_queue():
assert src.is_absolute(), 'expecting only absolute paths, not %r' % src
dst = pathlib.Path(dst).absolute()
try:
relpath = dst.relative_to(zippath)
# Don't bother trying to compress already-compressed files.
if src.suffix.lower() in STORE_ONLY:
compression = zipfile.ZIP_STORED
log.debug('ZIP %s -> %s (uncompressed)', src, relpath)
else:
compression = zipfile.ZIP_DEFLATED
log.debug('ZIP %s -> %s', src, relpath)
outzip.write(str(src), arcname=str(relpath), compress_type=compression)
if act == transfer.Action.MOVE:
self.delete_file(src)
except Exception:
# We have to catch exceptions in a broad way, as this is running in
# a separate thread, and exceptions won't otherwise be seen.
log.exception('Error transferring %s to %s', src, dst)
# Put the files to copy back into the queue, and abort. This allows
# the main thread to inspect the queue and see which files were not
# copied. The one we just failed (due to this exception) should also
# be reported there.
self.queue.put((src, dst, act))
return