RamEx-Flow / backend /app /utils /storage.py
Aye10032
feat(upload): support TAR and TGZ archives
0040215
Raw
History Blame Contribute Delete
3 kB
import json
import os
import shutil
import tarfile
import zipfile
from pathlib import Path
from typing import TYPE_CHECKING
from uuid import uuid4
from app.core.config import get_settings
if TYPE_CHECKING:
from fastapi import UploadFile
def project_workspace(project_id: str) -> Path:
path = get_settings().workspace_root / project_id
for child in ["input", "extracted", "scripts", "data", "plots", "results", "logs"]:
(path / child).mkdir(parents=True, exist_ok=True)
return path
def analysis_workspace(project_id: str, job_id: str) -> Path:
path = project_workspace(project_id) / "analysis" / job_id
for child in ["scripts", "results", "plots", "logs", "input"]:
(path / child).mkdir(parents=True, exist_ok=True)
return path
def remove_project_workspace(project_id: str) -> None:
path = get_settings().workspace_root / project_id
if path.exists():
shutil.rmtree(path)
def save_upload(project_id: str, upload: "UploadFile") -> Path:
workspace = project_workspace(project_id)
filename = (upload.filename or "upload.zip").lower()
suffix = next((ext for ext in (".tar.gz", ".tgz", ".tar", ".zip") if filename.endswith(ext)), ".zip")
safe_name = f"{uuid4()}{suffix}"
destination = workspace / "input" / safe_name
with destination.open("wb") as file_obj:
shutil.copyfileobj(upload.file, file_obj)
return destination
def extract_archive(project_id: str, archive_path: Path) -> Path:
workspace = project_workspace(project_id)
extract_dir = workspace / "extracted" / archive_path.name.replace(".", "_")
extract_dir.mkdir(parents=True, exist_ok=True)
if zipfile.is_zipfile(archive_path):
with zipfile.ZipFile(archive_path) as archive:
archive.extractall(extract_dir)
elif tarfile.is_tarfile(archive_path):
with tarfile.open(archive_path, "r:*") as archive:
archive.extractall(extract_dir, filter="data")
else:
raise ValueError("unsupported archive format")
return extract_dir
def infer_group_names(extract_dir: Path, group_index: int) -> list[str]:
groups: list[str] = []
seen: set[str] = set()
for root, _, files in os.walk(extract_dir):
for file_name in sorted(files):
if file_name.startswith("."):
continue
stem = Path(file_name).stem
parts = stem.split("_")
if len(parts) >= group_index:
group = parts[group_index - 1].strip()
else:
group = stem.strip()
if group and group not in seen:
seen.add(group)
groups.append(group)
return groups
def save_group_order(project_id: str, group_order: list[str]) -> Path:
workspace = project_workspace(project_id)
destination = workspace / "data" / "group_order.json"
destination.write_text(json.dumps(group_order, ensure_ascii=False, indent=2), encoding="utf-8")
return destination