mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2026-09-12 19:50:57 +05:00
80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
__package__ = "archivebox.misc"
|
|
|
|
# Post-bootstrap filesystem utilities (atomic_write, get_dir_size).
|
|
# Requires archivebox.config (uses get_config for permissions/output settings).
|
|
# Not safe to import pre-bootstrap.
|
|
|
|
import os
|
|
|
|
from json import dump
|
|
from pathlib import Path
|
|
|
|
from atomicwrites import atomic_write as lib_atomic_write
|
|
|
|
from archivebox.config.common import get_config
|
|
from archivebox.misc.util import enforce_types, ExtendedEncoder
|
|
|
|
|
|
@enforce_types
|
|
def atomic_write(path: Path | str, contents: dict | str | bytes, overwrite: bool = True, config=None, **config_kwargs) -> None:
|
|
"""Safe atomic write to filesystem by writing to temp file + atomic rename"""
|
|
|
|
mode = "wb+" if isinstance(contents, bytes) else "w"
|
|
encoding = None if isinstance(contents, bytes) else "utf-8" # enforce utf-8 on all text writes
|
|
|
|
try:
|
|
with lib_atomic_write(path, mode=mode, overwrite=overwrite, encoding=encoding) as f:
|
|
if isinstance(contents, dict):
|
|
dump(contents, f, indent=4, sort_keys=True, cls=ExtendedEncoder)
|
|
elif isinstance(contents, (bytes, str)):
|
|
f.write(contents)
|
|
except OSError as e:
|
|
config = config or get_config(**config_kwargs)
|
|
if config.ENFORCE_ATOMIC_WRITES:
|
|
print(f"[X] OSError: Failed to write {path} with fcntl.F_FULLFSYNC. ({e})")
|
|
print(
|
|
" You can store the archive/ subfolder on a hard drive or network share that doesn't support support synchronous writes,",
|
|
)
|
|
print(
|
|
" but the main folder containing the index.sqlite3 and ArchiveBox.conf files must be on a filesystem that supports FSYNC.",
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
# retry the write without forcing FSYNC (aka atomic mode)
|
|
with open(path, mode=mode, encoding=encoding) as f:
|
|
if isinstance(contents, dict):
|
|
dump(contents, f, indent=4, sort_keys=True, cls=ExtendedEncoder)
|
|
elif isinstance(contents, (bytes, str)):
|
|
f.write(contents)
|
|
|
|
# set file permissions
|
|
config = config or get_config(**config_kwargs)
|
|
os.chmod(path, int(config.OUTPUT_PERMISSIONS, base=8))
|
|
|
|
|
|
@enforce_types
|
|
def get_dir_size(path: str | Path, recursive: bool = True, pattern: str | None = None) -> tuple[int, int, int]:
|
|
"""get the total disk size of a given directory, optionally summing up
|
|
recursively and limiting to a given filter list
|
|
"""
|
|
num_bytes, num_dirs, num_files = 0, 0, 0
|
|
try:
|
|
for entry in os.scandir(path):
|
|
if (pattern is not None) and (pattern not in entry.path):
|
|
continue
|
|
if entry.is_dir(follow_symlinks=False):
|
|
if not recursive:
|
|
continue
|
|
num_dirs += 1
|
|
bytes_inside, dirs_inside, files_inside = get_dir_size(entry.path)
|
|
num_bytes += bytes_inside
|
|
num_dirs += dirs_inside
|
|
num_files += files_inside
|
|
else:
|
|
num_bytes += entry.stat(follow_symlinks=False).st_size
|
|
num_files += 1
|
|
except OSError:
|
|
# e.g. FileNameTooLong or other error while trying to read dir
|
|
pass
|
|
return num_bytes, num_dirs, num_files
|