102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from .schemas import JobPaths
|
|
|
|
|
|
_JOB_ID_RE = re.compile(r"^[0-9a-f]{32,}$")
|
|
|
|
|
|
class Storage:
|
|
def __init__(self, base_dir: Path) -> None:
|
|
self.base_dir = base_dir
|
|
self.base_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def prepare_job_dirs(self, job_id: str) -> JobPaths:
|
|
root = self.base_dir / job_id
|
|
input_dir = root / "input"
|
|
output_dir = root / "output"
|
|
input_dir.mkdir(parents=True, exist_ok=True)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
return JobPaths(
|
|
root=root,
|
|
input_dir=input_dir,
|
|
output_dir=output_dir,
|
|
mask_path=input_dir / "mask.png",
|
|
excel_path=input_dir / "names.xlsx",
|
|
config_path=root / "config.json",
|
|
)
|
|
|
|
def job_root(self, job_id: str) -> Path:
|
|
return self.base_dir / job_id
|
|
|
|
def job_dir_size(self, job_id: str) -> int:
|
|
root = self.job_root(job_id)
|
|
if not root.exists():
|
|
return 0
|
|
total = 0
|
|
for dirpath, _, filenames in os.walk(root):
|
|
for filename in filenames:
|
|
try:
|
|
total += os.path.getsize(os.path.join(dirpath, filename))
|
|
except OSError:
|
|
continue
|
|
return total
|
|
|
|
def job_dir_info(self, job_id: str) -> dict | None:
|
|
root = self.job_root(job_id)
|
|
if not root.is_dir():
|
|
return None
|
|
try:
|
|
mtime = root.stat().st_mtime
|
|
except OSError:
|
|
return None
|
|
return {
|
|
"job_id": job_id,
|
|
"path": str(root),
|
|
"size_bytes": self.job_dir_size(job_id),
|
|
"age_days": round(max(0.0, time.time() - mtime) / 86400.0, 2),
|
|
}
|
|
|
|
def list_job_ids(self) -> list[str]:
|
|
if not self.base_dir.exists():
|
|
return []
|
|
return [
|
|
item.name
|
|
for item in self.base_dir.iterdir()
|
|
if item.is_dir() and _JOB_ID_RE.match(item.name)
|
|
]
|
|
|
|
def stale_job_dirs(
|
|
self,
|
|
referenced_job_ids: set[str],
|
|
exclude_job_ids: set[str] | None = None,
|
|
max_age_days: float | None = None,
|
|
) -> list[dict]:
|
|
exclude = exclude_job_ids or set()
|
|
result: list[dict] = []
|
|
for job_id in self.list_job_ids():
|
|
if job_id in referenced_job_ids or job_id in exclude:
|
|
continue
|
|
info = self.job_dir_info(job_id)
|
|
if not info:
|
|
continue
|
|
if max_age_days is not None and info["age_days"] < max_age_days:
|
|
continue
|
|
result.append(info)
|
|
return sorted(result, key=lambda item: item["size_bytes"], reverse=True)
|
|
|
|
def remove_job_dir(self, job_id: str) -> int:
|
|
root = self.job_root(job_id)
|
|
if not root.exists():
|
|
return 0
|
|
size = self.job_dir_size(job_id)
|
|
shutil.rmtree(root, ignore_errors=True)
|
|
return size
|