Senko-san
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
"""LocalFileStorage — stores files on the local filesystem."""
|
|
|
|
import os
|
|
import shutil
|
|
from collections.abc import AsyncGenerator, AsyncIterator
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
import anyio
|
|
|
|
from app.domain.entities.storage import ObjectStat
|
|
from app.domain.errors import StorageError
|
|
|
|
_EXT_CONTENT_TYPE: dict[str, str] = {
|
|
"mp3": "audio/mpeg",
|
|
"flac": "audio/flac",
|
|
"m4a": "audio/mp4",
|
|
"aac": "audio/aac",
|
|
"ogg": "audio/ogg",
|
|
"opus": "audio/ogg",
|
|
"wav": "audio/wav",
|
|
"wma": "audio/x-ms-wma",
|
|
"aiff": "audio/aiff",
|
|
"aif": "audio/aiff",
|
|
}
|
|
|
|
|
|
class LocalFileStorage:
|
|
def __init__(self, media_path: Path) -> None:
|
|
self._media_path = media_path
|
|
|
|
async def save_file(self, key: str, src_path: Path) -> int:
|
|
dest = self._media_path / key
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
part = dest.with_suffix(dest.suffix + ".part")
|
|
shutil.copyfile(str(src_path), str(part))
|
|
os.replace(str(part), str(dest))
|
|
return dest.stat().st_size
|
|
|
|
async def open_range(
|
|
self, key: str, start: int, end: int | None
|
|
) -> tuple[AsyncIterator[bytes], int]:
|
|
path = self._media_path / key
|
|
if not path.exists():
|
|
raise StorageError(f"Object not found: {key}")
|
|
total_size = path.stat().st_size
|
|
|
|
_start = start
|
|
_end = end
|
|
_total_size = total_size
|
|
_path = path
|
|
|
|
async def _iter() -> AsyncGenerator[bytes]:
|
|
async with await anyio.open_file(_path, "rb") as f:
|
|
await f.seek(_start)
|
|
remaining = (_end - _start + 1) if _end is not None else (_total_size - _start)
|
|
while remaining > 0:
|
|
chunk: bytes = await f.read(min(65536, remaining))
|
|
if not chunk:
|
|
break
|
|
yield chunk
|
|
remaining -= len(chunk)
|
|
|
|
aiter: AsyncIterator[bytes] = _iter()
|
|
return aiter, total_size
|
|
|
|
async def stat(self, key: str) -> ObjectStat:
|
|
path = self._media_path / key
|
|
if not path.exists():
|
|
raise StorageError(f"Object not found: {key}")
|
|
st = path.stat()
|
|
ext = path.suffix.lower().lstrip(".")
|
|
return ObjectStat(size=st.st_size, content_type=_EXT_CONTENT_TYPE.get(ext))
|
|
|
|
async def exists(self, key: str) -> bool:
|
|
return (self._media_path / key).exists()
|
|
|
|
async def delete(self, key: str) -> None:
|
|
(self._media_path / key).unlink(missing_ok=True)
|
|
|
|
def as_local_path(self, key: str) -> AbstractAsyncContextManager[Path]:
|
|
return self._as_local_path_cm(key)
|
|
|
|
@asynccontextmanager
|
|
async def _as_local_path_cm(self, key: str) -> AsyncGenerator[Path]:
|
|
yield self._media_path / key
|