forked from wrenn/python-sdk
v0.1.0
This commit is contained in:
404
src/wrenn/files.py
Normal file
404
src/wrenn/files.py
Normal file
@ -0,0 +1,404 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
|
||||
import httpx
|
||||
|
||||
from wrenn.exceptions import WrennNotFoundError, 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},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
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},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
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')}"
|
||||
},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
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},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
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},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
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')}"
|
||||
},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user