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

update BAT module to 1.15 version

Browse Source
This commit is contained in:
pullusb
2023-01-10 11:41:55 +01:00
parent 85037c027a
commit fc32ebd453
34 changed files with 1878 additions and 1042 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pack/__init__.py
+115 -74
View File
@@ -28,6 +28,7 @@ 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__)
@@ -93,14 +94,16 @@ class Packer:
instance.
"""
def __init__(self,
bfile: pathlib.Path,
project: pathlib.Path,
target: str,
*,
noop=False,
compress=False,
relative_only=False) -> None:
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
@@ -110,7 +113,7 @@ class Packer:
self.relative_only = relative_only
self._aborted = threading.Event()
self._abort_lock = threading.RLock()
self._abort_reason = ''
self._abort_reason = ""
# Set this to a custom Callback() subclass instance before calling
# strategise() to receive progress reports.
@@ -119,15 +122,15 @@ class Packer:
self._exclude_globs = set() # type: typing.Set[str]
from blender_asset_tracer.cli import common
self._shorten = functools.partial(common.shorten, self.project)
self._shorten = functools.partial(shorten_path, self.project)
if noop:
log.warning('Running in no-op mode, only showing what will be done.')
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._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]
@@ -138,7 +141,7 @@ class Packer:
# Number of files we would copy, if not for --noop
self._file_count = 0
self._tmpdir = tempfile.TemporaryDirectory(prefix='bat-', suffix='-batpack')
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:
@@ -155,7 +158,7 @@ class Packer:
self._tscb.flush()
self._tmpdir.cleanup()
def __enter__(self) -> 'Packer':
def __enter__(self) -> "Packer":
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
@@ -177,7 +180,7 @@ class Packer:
self._progress_cb = new_progress_cb
self._tscb = progress.ThreadSafeCallback(self._progress_cb)
def abort(self, reason='') -> None:
def abort(self, reason="") -> None:
"""Aborts the current packing process.
Can be called from any thread. Aborts as soon as the running strategise
@@ -196,12 +199,12 @@ class Packer:
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')
log.error("A transfer error occurred")
reason = self._file_transferer.error_message()
elif not self._aborted.is_set():
return
log.warning('Aborting')
log.warning("Aborting")
self._tscb.flush()
self._progress_cb.pack_aborted(reason)
raise Aborted(reason)
@@ -212,8 +215,10 @@ class Packer:
Must be called before calling strategise().
"""
if self._actions:
raise RuntimeError('%s.exclude() must be called before strategise()' %
self.__class__.__qualname__)
raise RuntimeError(
"%s.exclude() must be called before strategise()"
% self.__class__.__qualname__
)
self._exclude_globs.update(globs)
def strategise(self) -> None:
@@ -236,7 +241,9 @@ class Packer:
# 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))
bfile_pp = self._target_path / bfile_path.relative_to(
bpathlib.make_absolute(self.project)
)
self._output_path = bfile_pp
self._progress_cb.pack_start()
@@ -251,11 +258,11 @@ class Packer:
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)
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)
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:
@@ -269,14 +276,22 @@ class Packer:
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)
def handle_missing_file():
log.warning("Missing file: %s", asset_path)
self.missing_files.add(asset_path)
self._progress_cb.missing_file(asset_path)
try:
for file_path in file_sequence.expand_sequence(asset_path):
if file_path.exists():
break
else:
# At least some file of a sequence must exist.
handle_missing_file()
return
except file_sequence.DoesNotExist:
# The asset path should point to something existing.
handle_missing_file()
return
# Handle this sequence as an asset.
@@ -291,7 +306,7 @@ class Packer:
# 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)
log.warning("Missing file: %s", asset_path)
self.missing_files.add(asset_path)
self._progress_cb.missing_file(asset_path)
return
@@ -315,11 +330,11 @@ class Packer:
act.usages.append(usage)
if needs_rewriting:
log.info('%s needs rewritten path to %s', bfile_path, usage.asset_path)
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)
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
@@ -331,7 +346,7 @@ class Packer:
assert isinstance(act, AssetAction)
relpath = bpathlib.strip_root(path)
act.new_path = pathlib.Path(self._target_path, '_outside_project', relpath)
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.
@@ -370,7 +385,7 @@ class Packer:
def execute(self) -> None:
"""Execute the strategy."""
assert self._actions, 'Run strategise() first'
assert self._actions, "Run strategise() first"
if not self.noop:
self._rewrite_paths()
@@ -408,7 +423,7 @@ class Packer:
This creates the BAT Pack but does not yet do any path rewriting.
"""
log.debug('Executing %d copy actions', len(self._actions))
log.debug("Executing %d copy actions", len(self._actions))
assert self._file_transferer is not None
@@ -418,12 +433,12 @@ class Packer:
self._copy_asset_and_deps(asset_path, action)
if self.noop:
log.info('Would copy %d files to %s', self._file_count, self.target)
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.')
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
@@ -460,23 +475,26 @@ class Packer:
# 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
assert bfile_pp is not None, \
f"Action {action.path_action.name} on {bfile_path} has no final path set, unable to process"
# 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_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)
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+')
bfile.copy_and_rebind(bfile_tp, mode="rb+")
for usage in action.rewrites:
self._check_aborted()
@@ -484,25 +502,35 @@ class Packer:
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)
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)
log.info(" - %s remained at %s", usage.asset_path, relpath)
continue
log.info(' - %s moved to %s', usage.asset_path, relpath)
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)
# Pointers can point to a non-existing data block, in which case
# either a SegmentationFault exception is thrown, or None is
# returned, based on the strict pointer mode set on the
# BlendFile class. Since this block was already meant to be
# rewritten, it was found before.
assert block is not None
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)
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)
log.debug(" - written %d bytes", written)
# BIG FAT ASSUMPTION that the filename (e.g. basename
# without path) does not change. This makes things much
@@ -510,10 +538,13 @@ class Packer:
# 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)
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)
log.debug(" - written %d bytes", written)
# Make sure we close the file, otherwise changes may not be
# flushed before it gets copied.
@@ -524,12 +555,13 @@ class Packer:
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):
if "*" not in str(asset_path) and "<UDIM>" not in asset_path.name:
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)
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:
@@ -541,7 +573,7 @@ class Packer:
# 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
assert "*" not in str(first_pp) or "*" in first_pp.name
packed_base_dir = first_pp.parent
for file_path in usage.files():
@@ -552,17 +584,16 @@ class Packer:
# 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):
def _send_to_target(
self, asset_path: pathlib.Path, target: pathlib.PurePath, may_move=False
):
if self.noop:
print('%s -> %s' % (asset_path, target))
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)
verb = "move" if may_move else "copy"
log.debug("Queueing %s of %s", verb, asset_path)
self._tscb.flush()
@@ -575,13 +606,23 @@ class Packer:
def _write_info_file(self):
"""Write a little text file with info at the top of the pack."""
infoname = 'pack-info.txt'
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)
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)
def shorten_path(cwd: pathlib.Path, somepath: pathlib.Path) -> pathlib.Path:
"""Return 'somepath' relative to CWD if possible."""
try:
return somepath.relative_to(cwd)
except ValueError:
return somepath
pack/filesystem.py
+33 -24
View File
@@ -81,9 +81,9 @@ class FileCopier(transfer.FileTransferer):
# 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)
log.debug("Error transferring %s to %s: %s", src, dst, ex)
else:
msg = 'Error transferring %s to %s' % (src, dst)
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
@@ -93,16 +93,16 @@ class FileCopier(transfer.FileTransferer):
self.queue.put((src, dst, act), timeout=1.0)
break
log.debug('All transfer threads queued')
log.debug("All transfer threads queued")
pool.close()
log.debug('Waiting for transfer threads to finish')
log.debug("Waiting for transfer threads to finish")
pool.join()
log.debug('All transfer threads finished')
log.debug("All transfer threads finished")
if self.files_transferred:
log.info('Transferred %d files', self.files_transferred)
log.info("Transferred %d files", self.files_transferred)
if self.files_skipped:
log.info('Skipped %d files', self.files_skipped)
log.info("Skipped %d files", self.files_skipped)
def _thread(self, src: pathlib.Path, dst: pathlib.Path, act: transfer.Action):
try:
@@ -111,7 +111,7 @@ class FileCopier(transfer.FileTransferer):
if self.has_error or self._abort.is_set():
raise AbortTransfer()
log.info('%s %s -> %s', act.name, src, dst)
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
@@ -121,9 +121,9 @@ class FileCopier(transfer.FileTransferer):
# 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)
log.debug("Error transferring %s to %s: %s", src, dst, ex)
else:
msg = 'Error transferring %s to %s' % (src, dst)
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
@@ -132,7 +132,9 @@ class FileCopier(transfer.FileTransferer):
# 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:
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():
@@ -142,20 +144,20 @@ class FileCopier(transfer.FileTransferer):
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)
log.info("SKIP %s; already exists", src)
if act == transfer.Action.MOVE:
log.debug('Deleting %s', src)
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"""
"""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))
"""Low-level file copy. dstpath needs to be a file and not a directory."""
shutil.copyfile(str(srcpath), str(dstpath))
def move(self, srcpath: pathlib.Path, dstpath: pathlib.Path):
s_stat = srcpath.stat()
@@ -171,19 +173,19 @@ class FileCopier(transfer.FileTransferer):
return
if (srcpath, dstpath) in self.already_copied:
log.debug('SKIP %s; already copied', srcpath)
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)
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)
log.debug("Copying %s -> %s", srcpath, dstpath)
self._copy(srcpath, dstpath)
self.already_copied.add((srcpath, dstpath))
@@ -191,8 +193,13 @@ class FileCopier(transfer.FileTransferer):
self.report_transferred(s_stat.st_size)
def copytree(self, src: pathlib.Path, dst: pathlib.Path,
symlinks=False, ignore_dangling_symlinks=False):
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:
@@ -204,7 +211,7 @@ class FileCopier(transfer.FileTransferer):
"""
if (src, dst) in self.already_copied:
log.debug('SKIP %s; already copied', src)
log.debug("SKIP %s; already copied", src)
return
if self.has_error or self._abort.is_set():
@@ -225,7 +232,9 @@ class FileCopier(transfer.FileTransferer):
# 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)
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:
@@ -250,7 +259,7 @@ class FileCopier(transfer.FileTransferer):
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:
if getattr(why, "winerror", None) is None:
errors.append((src, dst, str(why)))
if errors:
raise shutil.Error(errors)
pack/progress.py
+15 -10
View File
@@ -37,9 +37,11 @@ class Callback(blender_asset_tracer.trace.progress.Callback):
def pack_start(self) -> None:
"""Called when packing starts."""
def pack_done(self,
output_blendfile: pathlib.PurePath,
missing_files: typing.Set[pathlib.Path]) -> None:
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):
@@ -86,7 +88,7 @@ class ThreadSafeCallback(Callback):
"""
def __init__(self, wrapped: Callback) -> None:
self.log = log.getChild('ThreadSafeCallback')
self.log = log.getChild("ThreadSafeCallback")
self.wrapped = wrapped
# Thread-safe queue for passing progress reports on the main thread.
@@ -104,9 +106,11 @@ class ThreadSafeCallback(Callback):
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:
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):
@@ -135,8 +139,9 @@ class ThreadSafeCallback(Callback):
while True:
try:
call = self._reporting_queue.get(block=timeout is not None,
timeout=timeout)
call = self._reporting_queue.get(
block=timeout is not None, timeout=timeout
)
except queue.Empty:
return
@@ -145,4 +150,4 @@ class ThreadSafeCallback(Callback):
except Exception:
# Don't let the handling of one callback call
# block the entire flush process.
self.log.exception('Error calling %s', call)
self.log.exception("Error calling %s", call)
pack/s3.py
+40 -33
View File
@@ -32,17 +32,18 @@ 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)
log.debug("Computing MD5sum of %s", filepath)
hasher = hashlib.md5()
with filepath.open('rb') as infile:
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)
log.debug("MD5sum of %s is %s", filepath, md5)
return md5
@@ -63,20 +64,21 @@ class S3Packer(Packer):
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)
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)
self.client = self.session.client("s3", endpoint_url=endpoint)
def set_credentials(self,
endpoint: str,
access_key_id: str,
secret_access_key: str):
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)
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)
@@ -107,7 +109,7 @@ class S3Transferrer(transfer.FileTransferer):
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)
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
@@ -116,9 +118,9 @@ class S3Transferrer(transfer.FileTransferer):
return
if files_transferred:
log.info('Transferred %d files', files_transferred)
log.info("Transferred %d files", files_transferred)
if files_skipped:
log.info('Skipped %d files', 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.
@@ -135,25 +137,30 @@ class S3Transferrer(transfer.FileTransferer):
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)
log.debug(
"skipping %s, it already exists on the server with MD5 %s",
src,
existing_md5,
)
return False
log.info('Uploading %s', src)
log.info("Uploading %s", src)
try:
self.client.upload_file(str(src),
Bucket=bucket,
Key=key,
Callback=self.report_transferred,
ExtraArgs={'Metadata': {'md5': md5}})
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')
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]:
@@ -165,18 +172,18 @@ class S3Transferrer(transfer.FileTransferer):
"""
import botocore.exceptions
log.debug('Getting metadata of %s/%s', bucket, key)
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 = 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
if str(error_code) == "404":
return "", -1
raise ValueError("error response:" % ex.response) from None
try:
return info['Metadata']['md5'], info['ContentLength']
return info["Metadata"]["md5"], info["ContentLength"]
except KeyError:
return '', -1
return "", -1
pack/shaman/__init__.py
+40 -25
View File
@@ -17,7 +17,12 @@
# ***** END GPL LICENCE BLOCK *****
#
# (c) 2019, Blender Foundation - Sybren A. Stüvel
"""Shaman Client interface."""
"""Shaman Client interface.
Note that this supports the Shaman API of Flamenco Manager 2.x. Support for
Flamenco 3.x will be implemented in a new Flamenco Blender add-on, and not in
BAT itself.
"""
import logging
import os
import pathlib
@@ -38,13 +43,15 @@ 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:
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.
@@ -53,18 +60,20 @@ class ShamanPacker(bat_pack.Packer):
super().__init__(bfile, project, target, **kwargs)
self.checkout_id = checkout_id
self.shaman_endpoint = endpoint
self._checkout_location = ''
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')
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)
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
@@ -72,13 +81,17 @@ class ShamanPacker(bat_pack.Packer):
# 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)
return ShamanTransferrer(
auth_token, self.project, self.shaman_endpoint, self.checkout_id
)
def _make_target_path(self, target: str) -> pathlib.PurePath:
return pathlib.PurePosixPath('/')
return pathlib.PurePosixPath("/")
def _on_file_transfer_finished(self, *, file_transfer_completed: bool):
super()._on_file_transfer_finished(file_transfer_completed=file_transfer_completed)
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
@@ -104,7 +117,7 @@ class ShamanPacker(bat_pack.Packer):
try:
super().execute()
except requests.exceptions.ConnectionError as ex:
log.exception('Error communicating with Shaman')
log.exception("Error communicating with Shaman")
self.abort(str(ex))
self._check_aborted()
@@ -114,17 +127,19 @@ def parse_endpoint(shaman_url: str) -> typing.Tuple[str, str]:
urlparts = urllib.parse.urlparse(str(shaman_url))
if urlparts.scheme in {'shaman', 'shaman+https'}:
scheme = 'https'
elif urlparts.scheme == 'shaman+http':
scheme = 'http'
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)
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], '')
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
+30 -22
View File
@@ -30,7 +30,7 @@ from pathlib import Path
from . import time_tracker
CACHE_ROOT = Path().home() / '.cache/shaman-client/shasums'
CACHE_ROOT = Path().home() / ".cache/shaman-client/shasums"
MAX_CACHE_FILES_AGE_SECS = 3600 * 24 * 60 # 60 days
log = logging.getLogger(__name__)
@@ -53,7 +53,7 @@ def find_files(root: Path) -> typing.Iterable[Path]:
# Ignore hidden files/dirs; these can be things like '.svn' or '.git',
# which shouldn't be sent to Shaman.
if path.name.startswith('.'):
if path.name.startswith("."):
continue
if path.is_dir():
@@ -76,10 +76,10 @@ 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'):
log.debug("Computing checksum of %s", filepath)
with time_tracker.track_time(TimeInfo, "computing_checksums"):
hasher = hashlib.sha256()
with filepath.open('rb') as infile:
with filepath.open("rb") as infile:
while True:
block = infile.read(blocksize)
if not block:
@@ -98,7 +98,9 @@ def _cache_path(filepath: Path) -> Path:
# 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_key = (
base64.urlsafe_b64encode(encoded_path.encode(fs_encoding)).decode().rstrip("=")
)
cache_path = CACHE_ROOT / cache_key[:10] / cache_key[10:]
return cache_path
@@ -111,42 +113,44 @@ def compute_cached_checksum(filepath: Path) -> str:
skip the actual SHA256 computation.
"""
with time_tracker.track_time(TimeInfo, 'checksum_cache_handling'):
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:
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)
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):
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'):
with time_tracker.track_time(TimeInfo, "checksum_cache_handling"):
payload = {
'checksum': checksum,
'file_mtime': current_stat.st_mtime,
'file_size': current_stat.st_size,
"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:
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)
log.warning("Unable to write checksum cache file %s: %s", cache_path, ex)
return checksum
@@ -157,7 +161,7 @@ def cleanup_cache() -> None:
if not CACHE_ROOT.exists():
return
with time_tracker.track_time(TimeInfo, 'checksum_cache_handling'):
with time_tracker.track_time(TimeInfo, "checksum_cache_handling"):
queue = deque([CACHE_ROOT])
rmdir_queue = []
@@ -194,4 +198,8 @@ def cleanup_cache() -> None:
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)
log.info(
"Cache Cleanup: removed %d dirs and %d files",
num_removed_dirs,
num_removed_files,
)
pack/shaman/client.py
+14 -14
View File
@@ -37,14 +37,14 @@ class ShamanClient:
)
http_adapter = requests.adapters.HTTPAdapter(max_retries=retries)
self._session = requests.session()
self._session.mount('https://', http_adapter)
self._session.mount('http://', http_adapter)
self._session.mount("https://", http_adapter)
self._session.mount("http://", http_adapter)
if auth_token:
self._session.headers['Authorization'] = 'Bearer ' + auth_token
self._session.headers["Authorization"] = "Bearer " + auth_token
def request(self, method: str, url: str, **kwargs) -> requests.Response:
kwargs.setdefault('timeout', 300)
kwargs.setdefault("timeout", 300)
full_url = urllib.parse.urljoin(self._base_url, url)
return self._session.request(method, full_url, **kwargs)
@@ -56,8 +56,8 @@ class ShamanClient:
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', True)
return self.request('GET', url, **kwargs)
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.
@@ -67,8 +67,8 @@ class ShamanClient:
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', True)
return self.request('OPTIONS', url, **kwargs)
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.
@@ -78,8 +78,8 @@ class ShamanClient:
:rtype: requests.Response
"""
kwargs.setdefault('allow_redirects', False)
return self.request('HEAD', url, **kwargs)
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.
@@ -92,7 +92,7 @@ class ShamanClient:
:rtype: requests.Response
"""
return self.request('POST', url, data=data, json=json, **kwargs)
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.
@@ -104,7 +104,7 @@ class ShamanClient:
:rtype: requests.Response
"""
return self.request('PUT', url, data=data, **kwargs)
return self.request("PUT", url, data=data, **kwargs)
def patch(self, url, data=None, **kwargs):
r"""Sends a PATCH request. Returns :class:`Response` object.
@@ -116,7 +116,7 @@ class ShamanClient:
:rtype: requests.Response
"""
return self.request('PATCH', url, data=data, **kwargs)
return self.request("PATCH", url, data=data, **kwargs)
def delete(self, url, **kwargs):
r"""Sends a DELETE request. Returns :class:`Response` object.
@@ -126,4 +126,4 @@ class ShamanClient:
:rtype: requests.Response
"""
return self.request('DELETE', url, **kwargs)
return self.request("DELETE", url, **kwargs)
pack/shaman/transfer.py
+109 -64
View File
@@ -48,9 +48,15 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
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:
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
@@ -63,7 +69,7 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
# 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._checkout_location = ""
self.uploaded_files = 0
self.uploaded_bytes = 0
@@ -76,24 +82,32 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
# 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()
(
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))
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)
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)
to_upload = self._send_checkout_def_to_shaman(
definition_file, allowed_relpaths
)
if to_upload is None:
# An error has already been logged.
return
@@ -102,7 +116,7 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
break
# Send the files that still need to be sent.
self.log.info('Upload attempt %d', try_index + 1)
self.log.info("Upload attempt %d", try_index + 1)
failed_paths = self._upload_files(to_upload)
if not failed_paths:
break
@@ -113,11 +127,13 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
# 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)
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.log.info("All files uploaded succesfully")
self._request_checkout(definition_file)
# Delete the files that were supposed to be moved.
@@ -127,12 +143,13 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
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)
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]]:
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,
@@ -162,8 +179,8 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
filesize=filesize,
abspath=src,
)
line = '%s %s %s' % (checksum, filesize, relpath)
definition_lines.append(line.encode('utf8'))
line = "%s %s %s" % (checksum, filesize, relpath)
definition_lines.append(line.encode("utf8"))
relpaths.add(relpath)
if act == bat_transfer.Action.MOVE:
@@ -171,7 +188,7 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
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)
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
@@ -179,35 +196,39 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
# be reported there.
self.queue.put((src, dst, act))
self.error_set(msg)
return b'', set(), delete_when_done
return b"", set(), delete_when_done
cache.cleanup_cache()
return b'\n'.join(definition_lines), relpaths, delete_when_done
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]:
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)
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)
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)
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
msg = "Shaman requested path we did not intend to upload: %r" % path
self.log.error(msg)
self.error_set(msg)
return None
@@ -216,13 +237,13 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
to_upload.appendleft(path)
elif response == response_already_uploading:
to_upload.append(path)
elif response == 'ERROR':
msg = 'Error from Shaman: %s' % 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)
msg = "Unknown response from Shaman for path %r: %r" % (path, response)
self.log.error(msg)
self.error_set(msg)
return None
@@ -240,7 +261,9 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
def defer(some_path: str):
nonlocal to_upload
self.log.info(' %s deferred (already being uploaded by someone else)', some_path)
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.
@@ -251,35 +274,41 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
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.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))
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')
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)
self.log.info(" %s", path)
headers = {
'X-Shaman-Original-Filename': path,
"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))
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'
headers["X-Shaman-Can-Defer-Upload"] = "true"
url = 'files/%s/%d' % (fileinfo.checksum, fileinfo.filesize)
url = "files/%s/%d" % (fileinfo.checksum, fileinfo.filesize)
try:
with fileinfo.abspath.open('rb') as infile:
with fileinfo.abspath.open("rb") as infile:
resp = self.client.post(url, data=infile, headers=headers)
except requests.ConnectionError as ex:
@@ -290,7 +319,9 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
# connection.
defer(path)
else:
self.log.info(' %s could not be uploaded, might retry later: %s', path, ex)
self.log.info(
" %s could not be uploaded, might retry later: %s", path, ex
)
failed_paths.add(path)
continue
@@ -301,12 +332,15 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
if can_defer:
defer(path)
else:
self.log.info(' %s skipped (already existed on the server)', path)
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)
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
@@ -318,42 +352,53 @@ class ShamanTransferrer(bat_transfer.FileTransferer):
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)
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)
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')
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')
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'})
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)
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)
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 ""
return self._checkout_location
pack/transfer.py
+33 -17
View File
@@ -56,7 +56,7 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
def __init__(self) -> None:
super().__init__()
self.log = log.getChild('FileTransferer')
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
@@ -67,13 +67,15 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
# 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.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 = ''
self.__error_message = ""
# Instantiate a dummy progress callback so that we can call it
# without checking for None all the time.
@@ -87,8 +89,12 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
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'
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))
@@ -96,8 +102,12 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
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'
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))
@@ -107,7 +117,9 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
"""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)
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.
@@ -128,7 +140,8 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
assert files_remaining
raise FileTransferError(
"%d files couldn't be transferred" % len(files_remaining),
files_remaining)
files_remaining,
)
def _files_remaining(self) -> typing.List[pathlib.Path]:
"""Source files that were queued but not transferred."""
@@ -140,7 +153,7 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
def abort(self) -> None:
"""Abort the file transfer, immediately returns."""
log.info('Aborting')
log.info("Aborting")
self._abort.set()
def abort_and_join(self) -> None:
@@ -152,8 +165,11 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
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])
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."""
@@ -176,13 +192,13 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
if timeout:
run_until = time.time() + timeout
else:
run_until = float('inf')
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')
self.log.warning("Timeout while waiting for transfer to finish")
return
self.progress_cb.flush(timeout=0.5)
@@ -192,11 +208,11 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
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)
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)
log.warning("Unable to delete %s: %s", path, ex)
@property
def has_error(self) -> bool:
@@ -217,5 +233,5 @@ class FileTransferer(threading.Thread, metaclass=abc.ABCMeta):
"""Retrieve the error messsage, or an empty string if no error occurred."""
with self.__error_mutex:
if not self.__error.is_set():
return ''
return ""
return self.__error_message
pack/zipped.py
+9 -7
View File
@@ -30,7 +30,7 @@ from . import Packer, transfer
log = logging.getLogger(__name__)
# Suffixes to store uncompressed in the zip.
STORE_ONLY = {'.jpg', '.jpeg', '.exr'}
STORE_ONLY = {".jpg", ".jpeg", ".exr"}
class ZipPacker(Packer):
@@ -58,9 +58,9 @@ class ZipTransferrer(transfer.FileTransferer):
zippath = self.zippath.absolute()
with zipfile.ZipFile(str(zippath), 'w') as outzip:
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
assert src.is_absolute(), "expecting only absolute paths, not %r" % src
dst = pathlib.Path(dst).absolute()
try:
@@ -69,18 +69,20 @@ class ZipTransferrer(transfer.FileTransferer):
# 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)
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)
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)
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