79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
@dataclass
|
|
class ImportResult:
|
|
source: Path
|
|
target: Path
|
|
bytes: int
|
|
|
|
|
|
class ImportCancelled(Exception):
|
|
"""Raised when an import is cancelled at a safe copy boundary."""
|
|
|
|
|
|
class Importer:
|
|
def __init__(self, movies_root: Path, tv_root: Path):
|
|
self.movies_root = movies_root
|
|
self.tv_root = tv_root
|
|
|
|
def target_for(self, source: Path) -> Path:
|
|
target_root = self.tv_root if _looks_like_tv(source) else self.movies_root
|
|
return _unique_path(target_root / source.name)
|
|
|
|
def import_file(self, source: Path, should_cancel: Callable[[], bool] | None = None) -> ImportResult:
|
|
target = self.target_for(source)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
partial = target.with_name(target.name + ".partial")
|
|
try:
|
|
with source.open("rb") as src, partial.open("wb") as dst:
|
|
while True:
|
|
if should_cancel and should_cancel():
|
|
raise ImportCancelled("import cancelled")
|
|
chunk = src.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
dst.write(chunk)
|
|
dst.flush()
|
|
os.fsync(dst.fileno())
|
|
except ImportCancelled:
|
|
partial.unlink(missing_ok=True)
|
|
raise
|
|
if partial.stat().st_size != source.stat().st_size:
|
|
raise IOError("partial copy size mismatch")
|
|
partial.rename(target)
|
|
source.unlink()
|
|
_remove_empty_parents(source.parent)
|
|
return ImportResult(source=source, target=target, bytes=target.stat().st_size)
|
|
|
|
|
|
def _looks_like_tv(path: Path) -> bool:
|
|
text = str(path).lower()
|
|
return any(marker in text for marker in ("s01", "season", "episode"))
|
|
|
|
|
|
def _unique_path(path: Path) -> Path:
|
|
if not path.exists() and not path.with_name(path.name + ".partial").exists():
|
|
return path
|
|
stem, suffix = path.stem, path.suffix
|
|
for index in range(1, 10000):
|
|
candidate = path.with_name(f"{stem} ({index}){suffix}")
|
|
if not candidate.exists() and not candidate.with_name(candidate.name + ".partial").exists():
|
|
return candidate
|
|
raise RuntimeError(f"could not choose unique target for {path}")
|
|
|
|
|
|
def _remove_empty_parents(path: Path) -> None:
|
|
while True:
|
|
try:
|
|
path.rmdir()
|
|
except OSError:
|
|
return
|
|
path = path.parent
|