chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,337 @@
|
||||
import gzip
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from io import IOBase, TextIOBase
|
||||
from typing import IO, AsyncIterator, Dict, Iterator, Optional, Union, TypedDict
|
||||
|
||||
import httpx
|
||||
|
||||
from e2b.envd.filesystem import filesystem_pb2
|
||||
from e2b.exceptions import InvalidArgumentException
|
||||
from e2b.io_utils import agzip_iter, aiter_io_chunks, gzip_iter, iter_io_chunks
|
||||
|
||||
|
||||
class FileType(Enum):
|
||||
"""
|
||||
Enum representing the type of filesystem object.
|
||||
"""
|
||||
|
||||
FILE = "file"
|
||||
"""
|
||||
Filesystem object is a file.
|
||||
"""
|
||||
DIR = "dir"
|
||||
"""
|
||||
Filesystem object is a directory.
|
||||
"""
|
||||
|
||||
|
||||
def map_file_type(ft: filesystem_pb2.FileType):
|
||||
if ft == filesystem_pb2.FileType.FILE_TYPE_FILE:
|
||||
return FileType.FILE
|
||||
elif ft == filesystem_pb2.FileType.FILE_TYPE_DIRECTORY:
|
||||
return FileType.DIR
|
||||
|
||||
|
||||
def map_file_type_str(value: Optional[str]) -> Optional[FileType]:
|
||||
"""Map a `/files` API type string to `FileType`, `None` when unknown."""
|
||||
if value == FileType.FILE.value:
|
||||
return FileType.FILE
|
||||
elif value == FileType.DIR.value:
|
||||
return FileType.DIR
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class WriteInfo:
|
||||
"""
|
||||
Sandbox filesystem object information.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"""
|
||||
Name of the filesystem object.
|
||||
"""
|
||||
type: Optional[FileType]
|
||||
"""
|
||||
Type of the filesystem object.
|
||||
"""
|
||||
path: str
|
||||
"""
|
||||
Path to the filesystem object.
|
||||
"""
|
||||
metadata: Optional[Dict[str, str]] = field(default=None, kw_only=True)
|
||||
"""
|
||||
User-defined metadata stored on the file as `user.e2b.*` extended
|
||||
attributes. On writes this reflects the metadata supplied on upload; on
|
||||
reads (`get_info`, `list`, `rename`) it reflects any `user.e2b.*` xattr on
|
||||
the file, including ones set out-of-band. `None` when none is set.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: Dict) -> "WriteInfo":
|
||||
"""Build a `WriteInfo` from a `/files` upload response entry."""
|
||||
return cls(
|
||||
name=payload["name"],
|
||||
type=map_file_type_str(payload.get("type")),
|
||||
path=payload["path"],
|
||||
metadata=map_metadata(payload.get("metadata")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntryInfo(WriteInfo):
|
||||
"""
|
||||
Extended sandbox filesystem object information.
|
||||
"""
|
||||
|
||||
size: int
|
||||
"""
|
||||
Size of the filesystem object in bytes.
|
||||
"""
|
||||
mode: int
|
||||
"""
|
||||
File mode and permission bits.
|
||||
"""
|
||||
permissions: str
|
||||
"""
|
||||
String representation of file permissions (e.g. 'rwxr-xr-x').
|
||||
"""
|
||||
owner: str
|
||||
"""
|
||||
Owner of the filesystem object.
|
||||
"""
|
||||
group: str
|
||||
"""
|
||||
Group owner of the filesystem object.
|
||||
"""
|
||||
modified_time: datetime
|
||||
"""
|
||||
Last modification time of the filesystem object.
|
||||
"""
|
||||
symlink_target: Optional[str] = None
|
||||
"""
|
||||
Target of the symlink if the filesystem object is a symlink.
|
||||
If the filesystem object is not a symlink, this field is None.
|
||||
"""
|
||||
|
||||
|
||||
def map_entry_info(entry: filesystem_pb2.EntryInfo) -> EntryInfo:
|
||||
return EntryInfo(
|
||||
name=entry.name,
|
||||
type=map_file_type(entry.type),
|
||||
path=entry.path,
|
||||
size=entry.size,
|
||||
mode=entry.mode,
|
||||
permissions=entry.permissions,
|
||||
owner=entry.owner,
|
||||
group=entry.group,
|
||||
modified_time=entry.modified_time.ToDatetime(tzinfo=timezone.utc),
|
||||
# Optional, we can't directly access symlink_target otherwise it will be "" instead of None
|
||||
symlink_target=(
|
||||
entry.symlink_target if entry.HasField("symlink_target") else None
|
||||
),
|
||||
metadata=map_metadata(entry.metadata),
|
||||
)
|
||||
|
||||
|
||||
class WriteEntry(TypedDict):
|
||||
"""
|
||||
Contains path and data of the file to be written to the filesystem.
|
||||
"""
|
||||
|
||||
path: str
|
||||
data: Union[str, bytes, IO]
|
||||
|
||||
|
||||
class FileStreamReader(Iterator[bytes]):
|
||||
"""Iterator over a streamed file download.
|
||||
|
||||
Returned by ``Sandbox.files.read(format="stream")``. It owns the underlying
|
||||
HTTP response and releases its pooled connection as soon as the stream is
|
||||
fully consumed, an error is raised while reading (including the idle-read
|
||||
timeout, which raises ``httpx.ReadTimeout``), or the reader is closed.
|
||||
|
||||
There is no garbage-collection safety net, so always consume it fully, use
|
||||
it as a context manager, or call :meth:`close`::
|
||||
|
||||
with sandbox.files.read(path, format="stream") as stream:
|
||||
for chunk in stream:
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
self._iterator = response.iter_bytes()
|
||||
self._closed = False
|
||||
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
return self
|
||||
|
||||
def __next__(self) -> bytes:
|
||||
try:
|
||||
return next(self._iterator)
|
||||
except BaseException:
|
||||
# Covers normal end (StopIteration) and read errors alike.
|
||||
self.close()
|
||||
raise
|
||||
|
||||
def close(self) -> None:
|
||||
"""Release the underlying HTTP connection. Safe to call multiple times."""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
self._response.close()
|
||||
|
||||
def __enter__(self) -> "FileStreamReader":
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
class AsyncFileStreamReader(AsyncIterator[bytes]):
|
||||
"""Async iterator over a streamed file download.
|
||||
|
||||
Returned by ``AsyncSandbox.files.read(format="stream")``. It owns the
|
||||
underlying HTTP response and releases its pooled connection as soon as the
|
||||
stream is fully consumed, an error is raised while reading (including the
|
||||
idle-read timeout, which raises ``httpx.ReadTimeout``), or the reader is
|
||||
closed.
|
||||
|
||||
There is no garbage-collection safety net (releasing an async connection
|
||||
requires awaiting ``aclose()``, which a finalizer cannot do reliably), so
|
||||
always consume it fully, use it as an async context manager, or call
|
||||
:meth:`aclose`::
|
||||
|
||||
async with await sandbox.files.read(path, format="stream") as stream:
|
||||
async for chunk in stream:
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
self._iterator = response.aiter_bytes()
|
||||
self._closed = False
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
try:
|
||||
return await self._iterator.__anext__()
|
||||
except BaseException:
|
||||
# Covers normal end (StopAsyncIteration) and read errors alike.
|
||||
await self.aclose()
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Release the underlying HTTP connection. Safe to call multiple times."""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
await self._response.aclose()
|
||||
|
||||
async def __aenter__(self) -> "AsyncFileStreamReader":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info) -> None:
|
||||
await self.aclose()
|
||||
|
||||
|
||||
def _to_httpx_file(file_path: str, file_data: Union[str, bytes, IO]):
|
||||
"""Build an httpx multipart `("file", (name, data))` tuple for the upload."""
|
||||
if isinstance(file_data, (str, bytes)):
|
||||
return ("file", (file_path, file_data))
|
||||
elif isinstance(file_data, TextIOBase):
|
||||
return ("file", (file_path, file_data.read()))
|
||||
elif isinstance(file_data, IOBase):
|
||||
return ("file", (file_path, file_data))
|
||||
else:
|
||||
raise InvalidArgumentException(f"Unsupported data type for file {file_path}")
|
||||
|
||||
|
||||
def to_upload_body(
|
||||
data: Union[str, bytes, IO],
|
||||
use_gzip: bool = False,
|
||||
) -> Union[bytes, IO, Iterator[bytes]]:
|
||||
"""Prepare file data for upload, optionally gzip-compressed.
|
||||
|
||||
File-like objects are streamed in chunks instead of being buffered in
|
||||
memory.
|
||||
"""
|
||||
if isinstance(data, (str, bytes)):
|
||||
raw = data.encode("utf-8") if isinstance(data, str) else data
|
||||
return gzip.compress(raw) if use_gzip else raw
|
||||
elif isinstance(data, (TextIOBase, IOBase)):
|
||||
if use_gzip:
|
||||
return gzip_iter(iter_io_chunks(data))
|
||||
if isinstance(data, TextIOBase):
|
||||
# Text-mode IO yields str chunks—encode them while streaming.
|
||||
return iter_io_chunks(data)
|
||||
# httpx streams binary file-like objects in chunks without buffering.
|
||||
return data
|
||||
else:
|
||||
raise InvalidArgumentException(f"Unsupported data type: {type(data)}")
|
||||
|
||||
|
||||
def to_upload_body_async(
|
||||
data: Union[str, bytes, IO],
|
||||
use_gzip: bool = False,
|
||||
) -> Union[bytes, AsyncIterator[bytes]]:
|
||||
"""Prepare file data for upload with async httpx, optionally gzip-compressed.
|
||||
|
||||
File-like objects are streamed in chunks instead of being buffered in
|
||||
memory. Async httpx requires an async iterable for streamed request bodies.
|
||||
"""
|
||||
if isinstance(data, (str, bytes)):
|
||||
raw = data.encode("utf-8") if isinstance(data, str) else data
|
||||
return gzip.compress(raw) if use_gzip else raw
|
||||
elif isinstance(data, (TextIOBase, IOBase)):
|
||||
chunks = aiter_io_chunks(data)
|
||||
return agzip_iter(chunks) if use_gzip else chunks
|
||||
else:
|
||||
raise InvalidArgumentException(f"Unsupported data type: {type(data)}")
|
||||
|
||||
|
||||
METADATA_HEADER_PREFIX = "X-Metadata-"
|
||||
|
||||
# Metadata keys travel as `X-Metadata-<key>` HTTP header names, so they must be
|
||||
# valid header tokens (RFC 7230); values travel as header values, restricted to
|
||||
# printable US-ASCII.
|
||||
_METADATA_KEY_REGEX = re.compile(r"\A[A-Za-z0-9!#$%&'*+\-.^_`|~]+\Z")
|
||||
_METADATA_VALUE_REGEX = re.compile(r"\A[\x20-\x7e]*\Z")
|
||||
|
||||
|
||||
def validate_metadata(metadata: Optional[Dict[str, str]]) -> None:
|
||||
"""Validate metadata keys/values before they are sent as upload headers."""
|
||||
if not metadata:
|
||||
return
|
||||
for key, value in metadata.items():
|
||||
if not _METADATA_KEY_REGEX.match(key):
|
||||
raise InvalidArgumentException(
|
||||
f"Invalid metadata key {key!r}: keys must be non-empty and use only "
|
||||
"HTTP token characters (letters, digits and !#$%&'*+-.^_`|~)."
|
||||
)
|
||||
if not _METADATA_VALUE_REGEX.match(value):
|
||||
raise InvalidArgumentException(
|
||||
f"Invalid metadata value for key {key!r}: values must be printable US-ASCII."
|
||||
)
|
||||
|
||||
|
||||
def metadata_to_headers(
|
||||
metadata: Optional[Dict[str, str]],
|
||||
) -> Dict[str, str]:
|
||||
"""Translate user metadata into the `X-Metadata-*` upload headers envd reads."""
|
||||
if not metadata:
|
||||
return {}
|
||||
return {f"{METADATA_HEADER_PREFIX}{key}": value for key, value in metadata.items()}
|
||||
|
||||
|
||||
def map_metadata(metadata) -> Optional[Dict[str, str]]:
|
||||
"""Normalize a proto/HTTP metadata map: drop empties and return a plain dict or None."""
|
||||
if not metadata:
|
||||
return None
|
||||
return dict(metadata)
|
||||
Reference in New Issue
Block a user