forked from wrenn/python-sdk
405 lines
13 KiB
Python
405 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import AsyncIterator, Iterator
|
|
|
|
import httpx
|
|
|
|
from wrenn.exceptions import WrennNotFoundError, _raise_for_status, handle_response
|
|
from wrenn.models import FileEntry, ListDirResponse, MakeDirResponse
|
|
|
|
|
|
class Files:
|
|
"""Sync filesystem interface. Accessed via ``capsule.files``."""
|
|
|
|
def __init__(self, capsule_id: str, http: httpx.Client) -> None:
|
|
self._capsule_id = capsule_id
|
|
self._http = http
|
|
|
|
def read(self, path: str) -> str:
|
|
"""Read a file as a UTF-8 string.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Returns:
|
|
str: File contents decoded as UTF-8.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
return self.read_bytes(path).decode("utf-8", errors="replace")
|
|
|
|
def read_bytes(self, path: str) -> bytes:
|
|
"""Read a file as raw bytes.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Returns:
|
|
bytes: Raw file contents.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/read",
|
|
json={"path": path},
|
|
)
|
|
_raise_for_status(resp)
|
|
return resp.content
|
|
|
|
def write(self, path: str, data: str | bytes) -> None:
|
|
"""Write data to a file inside the capsule.
|
|
|
|
Creates parent directories if they do not exist.
|
|
|
|
Args:
|
|
path (str): Absolute destination path inside the capsule.
|
|
data (str | bytes): Content to write. Strings are UTF-8 encoded.
|
|
"""
|
|
if isinstance(data, str):
|
|
data = data.encode("utf-8")
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/write",
|
|
files={"file": ("upload", data)},
|
|
data={"path": path},
|
|
)
|
|
_raise_for_status(resp)
|
|
|
|
def list(self, path: str, depth: int = 1) -> list[FileEntry]:
|
|
"""List directory contents.
|
|
|
|
Args:
|
|
path (str): Absolute path to the directory inside the capsule.
|
|
depth (int): Recursion depth. ``1`` lists only immediate children.
|
|
Defaults to ``1``.
|
|
|
|
Returns:
|
|
list[FileEntry]: Entries in the directory.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/list",
|
|
json={"path": path, "depth": depth},
|
|
)
|
|
parsed = ListDirResponse.model_validate(handle_response(resp))
|
|
return parsed.entries or []
|
|
|
|
def exists(self, path: str) -> bool:
|
|
"""Check whether a path exists inside the capsule.
|
|
|
|
Args:
|
|
path (str): Absolute path to check.
|
|
|
|
Returns:
|
|
bool: ``True`` if the path exists.
|
|
"""
|
|
parent = os.path.dirname(path)
|
|
name = os.path.basename(path)
|
|
try:
|
|
entries = self.list(parent, depth=1)
|
|
except WrennNotFoundError:
|
|
return False
|
|
return any(e.name == name for e in entries)
|
|
|
|
def make_dir(self, path: str) -> FileEntry:
|
|
"""Create a directory (with parents). Idempotent.
|
|
|
|
Args:
|
|
path (str): Absolute path of the directory to create.
|
|
|
|
Returns:
|
|
FileEntry: The created (or already-existing) directory entry.
|
|
"""
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/mkdir",
|
|
json={"path": path},
|
|
)
|
|
if resp.status_code == 409:
|
|
try:
|
|
body = resp.json()
|
|
if body.get("error", {}).get("code") == "conflict":
|
|
parent = os.path.dirname(path)
|
|
name = os.path.basename(path)
|
|
for entry in self.list(parent, depth=1):
|
|
if entry.name == name:
|
|
return entry
|
|
except Exception:
|
|
pass
|
|
parsed = MakeDirResponse.model_validate(handle_response(resp))
|
|
if parsed.entry is None:
|
|
raise RuntimeError("mkdir response missing entry")
|
|
return parsed.entry
|
|
|
|
def remove(self, path: str) -> None:
|
|
"""Remove a file or directory recursively.
|
|
|
|
Args:
|
|
path (str): Absolute path to remove.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/remove",
|
|
json={"path": path},
|
|
)
|
|
handle_response(resp)
|
|
|
|
def upload_stream(self, path: str, stream: Iterator[bytes]) -> None:
|
|
"""Stream a large file into the capsule.
|
|
|
|
Prefer this over :meth:`write` when the file is too large to hold in
|
|
memory.
|
|
|
|
Args:
|
|
path (str): Absolute destination path inside the capsule.
|
|
stream (Iterator[bytes]): Iterable of byte chunks to upload.
|
|
"""
|
|
boundary = os.urandom(16).hex().encode("utf-8")
|
|
|
|
def _multipart() -> Iterator[bytes]:
|
|
yield b"--" + boundary + b"\r\n"
|
|
yield b'Content-Disposition: form-data; name="path"\r\n\r\n'
|
|
yield path.encode("utf-8") + b"\r\n"
|
|
yield b"--" + boundary + b"\r\n"
|
|
yield b'Content-Disposition: form-data; name="file"; filename="upload.bin"\r\n'
|
|
yield b"Content-Type: application/octet-stream\r\n\r\n"
|
|
for chunk in stream:
|
|
yield chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
|
yield b"\r\n--" + boundary + b"--\r\n"
|
|
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/stream/write",
|
|
content=_multipart(),
|
|
headers={
|
|
"Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}"
|
|
},
|
|
)
|
|
_raise_for_status(resp)
|
|
|
|
def download_stream(self, path: str) -> Iterator[bytes]:
|
|
"""Stream a large file out of the capsule.
|
|
|
|
Prefer this over :meth:`read_bytes` when the file is too large to hold
|
|
in memory.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Yields:
|
|
bytes: Successive byte chunks of the file.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
with self._http.stream(
|
|
"POST",
|
|
f"/v1/capsules/{self._capsule_id}/files/stream/read",
|
|
json={"path": path},
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
yield from resp.iter_bytes()
|
|
|
|
|
|
class AsyncFiles:
|
|
"""Async filesystem interface. Accessed via ``capsule.files``."""
|
|
|
|
def __init__(self, capsule_id: str, http: httpx.AsyncClient) -> None:
|
|
self._capsule_id = capsule_id
|
|
self._http = http
|
|
|
|
async def read(self, path: str) -> str:
|
|
"""Read a file as a UTF-8 string.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Returns:
|
|
str: File contents decoded as UTF-8.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
data = await self.read_bytes(path)
|
|
return data.decode("utf-8", errors="replace")
|
|
|
|
async def read_bytes(self, path: str) -> bytes:
|
|
"""Read a file as raw bytes.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Returns:
|
|
bytes: Raw file contents.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/read",
|
|
json={"path": path},
|
|
)
|
|
_raise_for_status(resp)
|
|
return resp.content
|
|
|
|
async def write(self, path: str, data: str | bytes) -> None:
|
|
"""Write data to a file inside the capsule.
|
|
|
|
Creates parent directories if they do not exist.
|
|
|
|
Args:
|
|
path (str): Absolute destination path inside the capsule.
|
|
data (str | bytes): Content to write. Strings are UTF-8 encoded.
|
|
"""
|
|
if isinstance(data, str):
|
|
data = data.encode("utf-8")
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/write",
|
|
files={"file": ("upload", data)},
|
|
data={"path": path},
|
|
)
|
|
_raise_for_status(resp)
|
|
|
|
async def list(self, path: str, depth: int = 1) -> list[FileEntry]:
|
|
"""List directory contents.
|
|
|
|
Args:
|
|
path (str): Absolute path to the directory inside the capsule.
|
|
depth (int): Recursion depth. ``1`` lists only immediate children.
|
|
Defaults to ``1``.
|
|
|
|
Returns:
|
|
list[FileEntry]: Entries in the directory.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/list",
|
|
json={"path": path, "depth": depth},
|
|
)
|
|
parsed = ListDirResponse.model_validate(handle_response(resp))
|
|
return parsed.entries or []
|
|
|
|
async def exists(self, path: str) -> bool:
|
|
"""Check whether a path exists inside the capsule.
|
|
|
|
Args:
|
|
path (str): Absolute path to check.
|
|
|
|
Returns:
|
|
bool: ``True`` if the path exists.
|
|
"""
|
|
parent = os.path.dirname(path)
|
|
name = os.path.basename(path)
|
|
try:
|
|
entries = await self.list(parent, depth=1)
|
|
except WrennNotFoundError:
|
|
return False
|
|
return any(e.name == name for e in entries)
|
|
|
|
async def make_dir(self, path: str) -> FileEntry:
|
|
"""Create a directory (with parents). Idempotent.
|
|
|
|
Args:
|
|
path (str): Absolute path of the directory to create.
|
|
|
|
Returns:
|
|
FileEntry: The created (or already-existing) directory entry.
|
|
"""
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/mkdir",
|
|
json={"path": path},
|
|
)
|
|
if resp.status_code == 409:
|
|
try:
|
|
body = resp.json()
|
|
if body.get("error", {}).get("code") == "conflict":
|
|
parent = os.path.dirname(path)
|
|
name = os.path.basename(path)
|
|
for entry in await self.list(parent, depth=1):
|
|
if entry.name == name:
|
|
return entry
|
|
except Exception:
|
|
pass
|
|
parsed = MakeDirResponse.model_validate(handle_response(resp))
|
|
if parsed.entry is None:
|
|
raise RuntimeError("mkdir response missing entry")
|
|
return parsed.entry
|
|
|
|
async def remove(self, path: str) -> None:
|
|
"""Remove a file or directory recursively.
|
|
|
|
Args:
|
|
path (str): Absolute path to remove.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/remove",
|
|
json={"path": path},
|
|
)
|
|
handle_response(resp)
|
|
|
|
async def upload_stream(self, path: str, stream: AsyncIterator[bytes]) -> None:
|
|
"""Stream a large file into the capsule.
|
|
|
|
Prefer this over :meth:`write` when the file is too large to hold in
|
|
memory.
|
|
|
|
Args:
|
|
path (str): Absolute destination path inside the capsule.
|
|
stream (AsyncIterator[bytes]): Async iterable of byte chunks to
|
|
upload.
|
|
"""
|
|
boundary = os.urandom(16).hex().encode("utf-8")
|
|
|
|
async def _multipart() -> AsyncIterator[bytes]:
|
|
yield b"--" + boundary + b"\r\n"
|
|
yield b'Content-Disposition: form-data; name="path"\r\n\r\n'
|
|
yield path.encode("utf-8") + b"\r\n"
|
|
yield b"--" + boundary + b"\r\n"
|
|
yield b'Content-Disposition: form-data; name="file"; filename="upload.bin"\r\n'
|
|
yield b"Content-Type: application/octet-stream\r\n\r\n"
|
|
async for chunk in stream:
|
|
yield chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
|
yield b"\r\n--" + boundary + b"--\r\n"
|
|
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/files/stream/write",
|
|
content=_multipart(),
|
|
headers={
|
|
"Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}"
|
|
},
|
|
)
|
|
_raise_for_status(resp)
|
|
|
|
async def download_stream(self, path: str) -> AsyncIterator[bytes]:
|
|
"""Stream a large file out of the capsule.
|
|
|
|
Prefer this over :meth:`read_bytes` when the file is too large to hold
|
|
in memory.
|
|
|
|
Args:
|
|
path (str): Absolute path to the file inside the capsule.
|
|
|
|
Yields:
|
|
bytes: Successive byte chunks of the file.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If the path does not exist.
|
|
"""
|
|
async with self._http.stream(
|
|
"POST",
|
|
f"/v1/capsules/{self._capsule_id}/files/stream/read",
|
|
json={"path": path},
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
async for chunk in resp.aiter_bytes():
|
|
yield chunk
|