2021-07-22 10:14:43 +02:00

81 lines
2.3 KiB
Python

"""shutil-like functionality while compressing blendfiles on the fly."""
import gzip
import logging
import pathlib
import shutil
log = logging.getLogger(__name__)
# Arbitrarily chosen block size, in bytes.
BLOCK_SIZE = 256 * 2 ** 10
def move(src: pathlib.Path, dest: pathlib.Path):
"""Move a file from src to dest, gzip-compressing if not compressed yet.
Only compresses files ending in .blend; others are moved as-is.
"""
my_log = log.getChild("move")
my_log.debug("Moving %s to %s", src, dest)
if src.suffix.lower() == ".blend":
_move_or_copy(src, dest, my_log, source_must_remain=False)
else:
shutil.move(str(src), str(dest))
def copy(src: pathlib.Path, dest: pathlib.Path):
"""Copy a file from src to dest, gzip-compressing if not compressed yet.
Only compresses files ending in .blend; others are copied as-is.
"""
my_log = log.getChild("copy")
my_log.debug("Copying %s to %s", src, dest)
if src.suffix.lower() == ".blend":
_move_or_copy(src, dest, my_log, source_must_remain=True)
else:
shutil.copy2(str(src), str(dest))
def _move_or_copy(
src: pathlib.Path,
dest: pathlib.Path,
my_log: logging.Logger,
*,
source_must_remain: bool
):
"""Either move or copy a file, gzip-compressing if not compressed yet.
:param src: File to copy/move.
:param dest: Path to copy/move to.
:source_must_remain: True to copy, False to move.
:my_log: Logger to use for logging.
"""
srcfile = src.open("rb")
try:
first_bytes = srcfile.read(2)
if first_bytes == b"\x1f\x8b":
# Already a gzipped file.
srcfile.close()
my_log.debug("Source file %s is GZipped already", src)
if source_must_remain:
shutil.copy2(str(src), str(dest))
else:
shutil.move(str(src), str(dest))
return
my_log.debug("Compressing %s on the fly while copying to %s", src, dest)
with gzip.open(str(dest), mode="wb") as destfile:
destfile.write(first_bytes)
shutil.copyfileobj(srcfile, destfile, BLOCK_SIZE)
srcfile.close()
if not source_must_remain:
my_log.debug("Deleting source file %s", src)
src.unlink()
finally:
if not srcfile.closed:
srcfile.close()