51 lines
No EOL
1.8 KiB
Python
51 lines
No EOL
1.8 KiB
Python
"""File handling for MAX surface."""
|
|
import os
|
|
import aiohttp
|
|
from pathlib import Path
|
|
|
|
|
|
class FileHandler:
|
|
def __init__(self, workspace_root: str):
|
|
self.workspace_root = workspace_root
|
|
|
|
def _make_unique_filename(self, directory: str, filename: str) -> str:
|
|
base = Path(filename).stem
|
|
ext = Path(filename).suffix
|
|
candidate = filename
|
|
counter = 1
|
|
while os.path.exists(os.path.join(directory, candidate)):
|
|
candidate = f"{base} ({counter}){ext}"
|
|
counter += 1
|
|
return candidate
|
|
|
|
async def download_attachment(
|
|
self,
|
|
download_url: str,
|
|
filename: str,
|
|
agent_workspace: str,
|
|
headers: dict = None,
|
|
) -> str:
|
|
full_dir = os.path.join(self.workspace_root, agent_workspace.strip("/"))
|
|
os.makedirs(full_dir, exist_ok=True)
|
|
|
|
unique_name = self._make_unique_filename(full_dir, filename)
|
|
filepath = os.path.join(full_dir, unique_name)
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(download_url, headers=headers) as resp:
|
|
resp.raise_for_status()
|
|
with open(filepath, "wb") as f:
|
|
f.write(await resp.read())
|
|
|
|
return unique_name
|
|
|
|
def read_outgoing_file(self, workspace_path: str, agent_workspace: str) -> bytes:
|
|
full_dir = os.path.join(self.workspace_root, agent_workspace.strip("/"))
|
|
filepath = os.path.join(full_dir, workspace_path.lstrip("/"))
|
|
with open(filepath, "rb") as f:
|
|
return f.read()
|
|
|
|
def file_exists(self, workspace_path: str, agent_workspace: str) -> bool:
|
|
full_dir = os.path.join(self.workspace_root, agent_workspace.strip("/"))
|
|
filepath = os.path.join(full_dir, workspace_path.lstrip("/"))
|
|
return os.path.exists(filepath) |