diff --git a/__init__.py b/__init__.py index 9bf1195..36c57f8 100755 --- a/__init__.py +++ b/__init__.py @@ -79,12 +79,20 @@ class ExportBatPack(Operator, ExportHelper): self.report({'INFO'},'Executing ZipPacker ...') - with zipped.ZipPacker( - Path(bpy.data.filepath), - Path(bpy.data.filepath).parent, - str(self.filepath)) as packer: - packer.strategise() - packer.execute() + try: + with zipped.ZipPacker( + Path(bpy.data.filepath), + Path(bpy.data.filepath).parent, + str(self.filepath)) as packer: + packer.strategise() + packer.execute() + except EnvironmentError as err: + self.report({'ERROR'}, "BAT packing failed: %s\n" + "The zstandard Python module may be missing or incompatible." % err) + return {'CANCELLED'} + except Exception as err: + self.report({'ERROR'}, "BAT packing failed: %s" % err) + return {'CANCELLED'} self.report({'INFO'},'Packing successful !') with zipfile.ZipFile(str(self.filepath)) as inzip: @@ -216,7 +224,15 @@ class BAT_OT_export_zip(Operator, ExportHelper): current_file = Path(bpy.data.filepath) links.append(str(current_file)) - file_link = list(deps(current_file)) + try: + file_link = list(deps(current_file)) + except EnvironmentError as err: + self.report({'ERROR'}, "BAT dependency tracing failed: %s\n" + "The zstandard Python module may be missing or incompatible." % err) + return {'CANCELLED'} + except Exception as err: + self.report({'ERROR'}, "BAT dependency tracing failed: %s" % err) + return {'CANCELLED'} for l in file_link: if Path(l.abspath).exists() == False: continue diff --git a/blendfile/magic_compression.py b/blendfile/magic_compression.py index 51a1cc4..6e98ef4 100755 --- a/blendfile/magic_compression.py +++ b/blendfile/magic_compression.py @@ -87,22 +87,24 @@ def open(path: pathlib.Path, mode: str, buffer_size: int) -> DecompressedFileInf ) log.debug("%s-compressed blendfile detected: %s", compression.name, path) + print("BAT: %s-compressed blendfile: %s" % (compression.name, path)) # Decompress to a temporary file. tmpfile = tempfile.NamedTemporaryFile() fileobj.seek(0, os.SEEK_SET) - decompressor = _decompressor(fileobj, mode, compression) + _decompress_to_file(fileobj, tmpfile, mode, compression, buffer_size) - with decompressor as compressed_file: - magic = compressed_file.read(len(BLENDFILE_MAGIC)) - if magic != BLENDFILE_MAGIC: - raise exceptions.BlendFileError("Compressed file is not a blend file", path) - - data = magic - while data: - tmpfile.write(data) - data = compressed_file.read(buffer_size) + # Verify the decompressed content is a valid blend file. + tmpfile.seek(0, os.SEEK_SET) + header_bytes = tmpfile.read(12) + log.debug("Decompressed header (%s, %s): %r", compression.name, path, header_bytes) + print("BAT: Decompressed header: %r" % header_bytes) + if not header_bytes.startswith(BLENDFILE_MAGIC): + raise exceptions.BlendFileError( + "Decompressed file is not a blend file (header: %r)" % header_bytes, path + ) + tmpfile.seek(0, os.SEEK_SET) # Further interaction should be done with the uncompressed file. fileobj.close() @@ -113,6 +115,37 @@ def open(path: pathlib.Path, mode: str, buffer_size: int) -> DecompressedFileInf ) +def _decompress_to_file( + fileobj: typing.IO[bytes], + tmpfile: typing.IO[bytes], + mode: str, + compression: Compression, + buffer_size: int, +) -> None: + """Decompress fileobj into tmpfile.""" + + if compression == Compression.GZIP: + with gzip.GzipFile(fileobj=fileobj, mode=mode) as gz: + while True: + data = gz.read(buffer_size) + if not data: + break + tmpfile.write(data) + + elif compression == Compression.ZSTD: + if not has_zstandard: + raise EnvironmentError( + "File is compressed with ZStandard, install the `zstandard` module to support this." + ) + dctx = zstandard.ZstdDecompressor() + dctx.copy_stream(fileobj, tmpfile) + + else: + raise ValueError("Unsupported compression type: %s" % compression) + + tmpfile.flush() + + def find_compression_type(fileobj: typing.IO[bytes]) -> Compression: fileobj.seek(0, os.SEEK_SET) @@ -150,20 +183,3 @@ def _matches_magic(value: bytes, magic: bytes) -> bool: return value[: len(magic)] == magic -def _decompressor( - fileobj: typing.IO[bytes], mode: str, compression: Compression -) -> typing.IO[bytes]: - if compression == Compression.GZIP: - decompressor = gzip.GzipFile(fileobj=fileobj, mode=mode) - return typing.cast(typing.IO[bytes], decompressor) - - if compression == Compression.ZSTD: - if not has_zstandard: - # The required module was not loaded, raise an exception about this. - raise EnvironmentError( - "File is compressed with ZStandard, install the `zstandard` module to support this." - ) - dctx = zstandard.ZstdDecompressor() - return dctx.stream_reader(fileobj) - - raise ValueError("Unsupported compression type: %s" % compression) diff --git a/preferences.py b/preferences.py index 1d3fa12..aa49f9a 100755 --- a/preferences.py +++ b/preferences.py @@ -11,7 +11,7 @@ def get_addon_prefs(): return prefs.addons[__package__].preferences class myaddonPrefs(bpy.types.AddonPreferences): - bl_idname = __name__.split('.')[0] + bl_idname = __package__ root_default : bpy.props.StringProperty( name='Default Root', diff --git a/trace/file2blocks.py b/trace/file2blocks.py index af5ff2e..c1428ef 100755 --- a/trace/file2blocks.py +++ b/trace/file2blocks.py @@ -133,7 +133,11 @@ class BlockIterator: continue log.debug("Expanding %d blocks in %s", len(idblocks), lib_path) - libfile = self.open_blendfile(lib_path) + try: + libfile = self.open_blendfile(lib_path) + except Exception: + log.warning("Failed to open library %s, skipping", lib_path, exc_info=True) + continue yield from self.iter_blocks(libfile, idblocks) def _queue_all_blocks(self, bfile: blendfile.BlendFile):