# Copyright 2022 The PyAthena authors
#
# Licensed under the MIT License.
# See LICENSE or https://opensource.org/licenses/MIT.
#
# SPDX-License-Identifier: MIT
"""Asynchronous fsspec filesystem for Amazon S3 built on ``S3FileSystem``."""
from __future__ import annotations
import asyncio
import logging
import mimetypes
import os
from multiprocessing import cpu_count
from typing import TYPE_CHECKING, Any, cast
from fsspec.asyn import AsyncFileSystem, sync
from fsspec.callbacks import _DEFAULT_CALLBACK
from fsspec.core import get_compression
from fsspec.implementations.local import LocalFileSystem, make_path_posix
from fsspec.utils import check_contained
from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem
from pyathena.filesystem.s3_core import S3Core
from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor
from pyathena.filesystem.s3_object import (
S3CompleteMultipartUpload,
S3Metadata,
S3MultipartUpload,
S3MultipartUploadPart,
S3Object,
S3ObjectType,
S3ObjectVersion,
)
from pyathena.filesystem.s3_path import S3Path
from pyathena.filesystem.s3_path_pairing import S3PathPairing
if TYPE_CHECKING:
from datetime import datetime
from pyathena.connection import Connection
_logger = logging.getLogger(__name__)
# The running cleanups of cancelled multipart copies.
_cleanup_tasks: set[asyncio.Future[None]] = set()
[docs]
class AioS3FileSystem(AsyncFileSystem):
"""An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
This class wraps ``S3FileSystem`` to provide native asyncio support. Instead of
using ``ThreadPoolExecutor`` for parallel operations, it uses ``asyncio.gather``
with ``asyncio.to_thread`` for natural integration with the asyncio event loop.
The implementation uses composition: an internal ``S3FileSystem`` instance handles
all boto3 calls, while this class delegates to it via ``asyncio.to_thread()``.
This avoids diamond inheritance issues and keeps all boto3 logic in one place.
File handles created by ``_open`` use ``S3AioExecutor`` so that parallel
operations (range reads, multipart uploads) are dispatched through the event
loop with ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file.
An instance created with ``asynchronous=True`` has no event loop of its own,
so its file handles use a ``ThreadPoolExecutor``.
Attributes:
_sync_fs: The internal synchronous S3FileSystem instance.
Example:
>>> from pyathena.filesystem.s3_async import AioS3FileSystem
>>> fs = AioS3FileSystem(asynchronous=True)
>>>
>>> # Use in async context
>>> files = await fs._ls('s3://my-bucket/data/')
>>>
>>> # Sync wrappers (auto-generated by fsspec) need an instance created
>>> # without asynchronous=True; they block the caller until done
>>> files = AioS3FileSystem().ls('s3://my-bucket/data/')
"""
protocol = ("s3", "s3a")
mirror_sync_methods = True
async_impl = True
_extra_tokenize_attributes = ("default_block_size",)
[docs]
def __init__(
self,
connection: Connection[Any] | None = None,
default_block_size: int | None = None,
default_cache_type: str | None = None,
max_workers: int = (cpu_count() or 1) * 5,
s3_additional_kwargs: dict[str, Any] | None = None,
allow_bucket_creation: bool = False,
allow_bucket_deletion: bool = False,
version_aware: bool = False,
asynchronous: bool = False,
loop: Any | None = None,
batch_size: int | None = None,
**kwargs,
) -> None:
"""Initialize the filesystem and its internal ``S3FileSystem``.
Args:
connection: Passed to the internal ``S3FileSystem``.
default_block_size: Passed to the internal ``S3FileSystem``.
default_cache_type: Passed to the internal ``S3FileSystem``.
max_workers: Passed to the internal ``S3FileSystem``.
s3_additional_kwargs: Passed to the internal ``S3FileSystem``.
allow_bucket_creation: Passed to the internal ``S3FileSystem``.
allow_bucket_deletion: Passed to the internal ``S3FileSystem``.
version_aware: Passed to the internal ``S3FileSystem``.
asynchronous: Passed to ``fsspec.asyn.AsyncFileSystem``.
loop: Passed to ``fsspec.asyn.AsyncFileSystem``.
batch_size: Passed to ``fsspec.asyn.AsyncFileSystem``.
**kwargs: Passed to both ``fsspec.asyn.AsyncFileSystem`` and the
internal ``S3FileSystem``.
"""
super().__init__(
asynchronous=asynchronous,
loop=loop,
batch_size=batch_size,
**kwargs,
)
self._sync_fs = S3FileSystem(
connection=connection,
default_block_size=default_block_size,
default_cache_type=default_cache_type,
max_workers=max_workers,
s3_additional_kwargs=s3_additional_kwargs,
allow_bucket_creation=allow_bucket_creation,
allow_bucket_deletion=allow_bucket_deletion,
version_aware=version_aware,
# fsspec caches the AioS3FileSystem itself when caching is wanted.
skip_instance_cache=True,
**kwargs,
)
# Share dircache for cache coherence between async and sync instances
self.dircache = self._sync_fs.dircache
@property
def core(self) -> S3Core:
"""The typed S3 operations of the wrapped ``S3FileSystem``."""
return self._sync_fs.core
async def _info(self, path: str, **kwargs) -> S3Object:
return await asyncio.to_thread(self._sync_fs.info, path, **kwargs)
async def _ls(self, path: str, detail: bool = False, **kwargs) -> list[S3Object] | list[str]:
return await asyncio.to_thread(self._sync_fs.ls, path, detail=detail, **kwargs)
async def _cat_file(
self, path: str, start: int | None = None, end: int | None = None, **kwargs
) -> bytes:
return await asyncio.to_thread(self._sync_fs.cat_file, path, start=start, end=end, **kwargs)
async def _exists(self, path: str, **kwargs) -> bool:
return await asyncio.to_thread(self._sync_fs.exists, path, **kwargs)
async def _rm_file(self, path: str, **kwargs) -> None:
await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs)
async def _pipe_file(
self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs
) -> None:
if self._intrans:
# The transaction belongs to this filesystem, not to the internal
# S3FileSystem, so write through open() to defer the commit to it.
await asyncio.to_thread(self._pipe_file_in_transaction, path, value, mode, **kwargs)
return
await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
def _pipe_file_in_transaction(
self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs
) -> None:
"""Write bytes into the path as a file of this filesystem's transaction.
Args:
path: S3 path (s3://bucket/key) to write to.
value: The bytes to write.
mode: "overwrite" or "create". With "create", the file is
opened in ``xb`` mode: raise FileExistsError when the object
already exists, including one created before the
transaction is committed, which is not replaced.
**kwargs: Additional parameters passed to ``open()``, except
``compression``, with which the value is compressed before
it is written, as in :meth:`S3FileSystem.pipe_file`.
Raises:
FileExistsError: If the mode is "create" and the path already
exists.
ValueError: If the compression is not supported, or if the data
takes more than ``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks.
"""
# See S3FileSystem.pipe_file.
compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None))
if compression is not None:
value = CompressedBuffer.compress(value, compression)
block_size = kwargs.get("block_size") or self._sync_fs.default_block_size
# The size in bytes; the length of a memoryview counts its items.
self.core.check_multipart_upload_size(memoryview(value).nbytes, block_size)
self.open(path, "xb" if mode == "create" else "wb", **kwargs)._write_and_close(value)
async def _put_file(
self,
lpath: str,
rpath: str,
callback=_DEFAULT_CALLBACK,
mode: str = "overwrite",
**kwargs,
) -> None:
if self._intrans:
# See _pipe_file.
await asyncio.to_thread(
self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs
)
return
await asyncio.to_thread(
self._sync_fs.put_file, lpath, rpath, callback=callback, mode=mode, **kwargs
)
def _put_file_in_transaction(
self, lpath: str, rpath: str, callback, mode: str, **kwargs
) -> None:
"""Upload a local file as a file of this filesystem's transaction.
Mirrors :meth:`S3FileSystem.put_file`, but writes through ``open()``
of this filesystem.
Args:
lpath: Local file path to upload.
rpath: S3 destination path (s3://bucket/key).
callback: Progress callback for tracking upload progress.
mode: "overwrite" or "create". With "create", the file is
opened in ``xb`` mode: raise FileExistsError when the object
already exists, including one created before the
transaction is committed, which is not replaced.
**kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
The ``block_size``, ``max_workers``, and ``s3_additional_kwargs``
parameters of ``open()`` are also accepted.
Raises:
FileExistsError: If the mode is "create" and the path already
exists.
ValueError: If the file takes more than
``S3Core.MULTIPART_UPLOAD_MAX_PARTS`` blocks.
"""
if os.path.isdir(lpath):
return
if not S3Path.parse(rpath).key:
return
size = os.path.getsize(lpath)
block_size = kwargs.pop("block_size", None) or self._sync_fs.default_block_size
max_workers = kwargs.pop("max_workers", self._sync_fs.max_workers)
# The other parameters are S3 request parameters, as in pipe_file().
s3_additional_kwargs = {**kwargs.pop("s3_additional_kwargs", {}), **kwargs}
self.core.check_multipart_upload_size(size, block_size)
callback.set_size(size)
if "ContentType" not in {**self._sync_fs.s3_additional_kwargs, **s3_additional_kwargs}:
content_type, _ = mimetypes.guess_type(lpath)
if content_type is not None:
s3_additional_kwargs["ContentType"] = content_type
# See S3FileSystem.put_file.
with open(lpath, "rb") as local:
self.open(
rpath,
"xb" if mode == "create" else "wb",
block_size=block_size,
max_workers=max_workers,
s3_additional_kwargs=s3_additional_kwargs,
)._write_file_and_close(local, callback)
self.invalidate_cache(rpath)
async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs)
async def _mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None:
await asyncio.to_thread(self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs)
async def _makedirs(self, path: str, exist_ok: bool = False) -> None:
await asyncio.to_thread(self._sync_fs.makedirs, path, exist_ok=exist_ok)
async def _rm(
self,
path: str | list[str],
recursive: bool = False,
batch_size: int | None = None,
maxdepth: int | None = None,
**kwargs,
) -> None:
"""Delete objects with DeleteObjects requests.
See :meth:`S3FileSystem.rm`. The requests run in parallel with
``asyncio.gather`` and ``asyncio.to_thread``.
Args:
path: S3 path (s3://bucket/key) or list of paths to delete.
recursive: Whether to delete all objects below the paths.
batch_size: Accepted for fsspec compatibility; not used.
maxdepth: Maximum depth to expand when ``recursive`` is True.
**kwargs: Additional parameters passed to the DeleteObjects API.
``Quiet`` (default True) sets the quiet mode of the requests.
Raises:
ValueError: If a path is a bucket.
OSError: If S3 could not delete some of the objects.
"""
paths = await asyncio.to_thread(
self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth
)
await self._delete_objects(paths, **kwargs)
async def _delete_objects(self, paths: list[str], **kwargs) -> None:
"""Delete objects with DeleteObjects requests run with ``asyncio.gather``.
See :meth:`S3FileSystem._delete_objects`.
Args:
paths: Paths of the objects to delete. Bucket paths are skipped.
**kwargs: Additional parameters passed to the DeleteObjects API.
``Quiet`` (default True) sets the quiet mode of the requests.
Raises:
TypeError: If kwargs has ``Bucket`` or ``Delete``.
OSError: If S3 could not delete some of the objects.
"""
batches, params = self._sync_fs._delete_batches(paths, **kwargs)
results = await asyncio.gather(
*[asyncio.to_thread(self._sync_fs._delete_batch, batch, **params) for batch in batches],
return_exceptions=True,
)
self._sync_fs._raise_delete_errors(results)
async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None:
"""Move files from one S3 location to another.
See :meth:`S3FileSystem.mv`. The copies and the deletions run in
parallel with ``asyncio.gather``.
Args:
path1: Source S3 path, glob pattern, or list of paths.
path2: Destination S3 path, or list of paths when ``path1`` is a
list.
recursive: Whether to move the directories with their contents.
maxdepth: Maximum depth of a recursive move.
**kwargs: Additional S3 copy parameters, as for ``_cp_file()``.
Raises:
ValueError: If two sources have the same destination, or a
destination is another source, including one left in place,
which is checked before anything is copied. A directory with
no object at its key, which is not copied, does not conflict.
OSError: If bucket versioning cannot be read.
"""
if path1 == path2:
return
pairs = await asyncio.to_thread(
self._sync_fs._move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth
)
# Every copy finishes before a failure is raised, as in fsspec's
# _copy(), and nothing is deleted after a failure.
results = await asyncio.gather(
*[self._copy_file(p1, p2, **kwargs) for p1, p2 in pairs], return_exceptions=True
)
for result in results:
if isinstance(result, BaseException):
raise result
await self._delete_objects(
[p1 for (p1, _), copied in zip(pairs, results, strict=True) if copied]
)
[docs]
def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None:
"""Move files from one S3 location to another.
Runs :meth:`_mv`; fsspec does not provide the synchronous wrapper of
``_mv``.
Args:
path1: Source S3 path, glob pattern, or list of paths.
path2: Destination S3 path, or list of paths when ``path1`` is a
list.
recursive: Whether to move the directories with their contents.
maxdepth: Maximum depth of a recursive move.
**kwargs: Additional S3 copy parameters.
Raises:
ValueError: If the move has conflicting paths.
OSError: If bucket versioning cannot be read.
"""
sync(self.loop, self._mv, path1, path2, recursive=recursive, maxdepth=maxdepth, **kwargs)
async def _copy(
self,
path1,
path2,
recursive=False,
on_error=None,
maxdepth=None,
batch_size=None,
**kwargs,
) -> None:
"""Copy files within S3.
See :meth:`S3FileSystem.copy`. The copies run as in fsspec's
``_copy()``.
Args:
path1: Source S3 path, glob pattern, or list of them.
path2: Destination S3 path, or list of paths when ``path1`` is a
list.
recursive: Whether to copy the directories with their contents.
on_error: ``"raise"`` or ``"ignore"`` for a missing source.
maxdepth: Maximum depth of a recursive copy.
batch_size: Number of copies to run at the same time.
**kwargs: Additional S3 copy parameters passed to ``_cp_file()``.
"""
sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1
if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources):
pairs = await asyncio.to_thread(
self._sync_fs._copy_pairs,
S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth),
)
if not pairs:
return
path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs]
await super()._copy(
path1,
path2,
recursive=recursive,
on_error=on_error,
maxdepth=maxdepth,
batch_size=batch_size,
**kwargs,
)
async def _get(
self, rpath, lpath, recursive=False, callback=_DEFAULT_CALLBACK, maxdepth=None, **kwargs
) -> None:
"""Copy files from S3 to the local filesystem.
See :meth:`S3FileSystem.get`. The downloads run as in fsspec's
``_get()``.
Args:
rpath: Source S3 path, glob pattern, or list of them.
lpath: Local destination path, or list of paths when ``rpath`` is
a list.
recursive: Whether to copy the directories with their contents.
callback: Progress callback.
maxdepth: Maximum depth of a recursive copy.
**kwargs: Additional parameters passed to ``_get_file()``.
Raises:
ValueError: If a source with a version ID is paired, and a
destination lies outside ``lpath``.
"""
sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath
if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources):
root = make_path_posix(lpath)
pairs = await asyncio.to_thread(
self._sync_fs._copy_pairs,
S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth),
isdir=LocalFileSystem().isdir,
)
rpath, lpath = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs]
check_contained(root, lpath)
if not rpath:
return
await super()._get(
rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs
)
async def _cp_file(self, path1: str, path2: str, **kwargs) -> None:
"""Copy an S3 object, using async parallel multipart upload for large files.
A directory ``path1``, which recursive ``copy()`` passes along with
the files under it, is skipped.
Args:
path1: Source S3 path (s3://bucket/key).
path2: Destination S3 path (s3://bucket/key).
**kwargs: Additional S3 copy parameters.
Raises:
ValueError: If trying to copy to a versioned file or copy buckets.
"""
await self._copy_file(path1, path2, **kwargs)
async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool:
"""Copy an S3 object as :meth:`_cp_file` does.
Args:
path1: Source S3 path (s3://bucket/key).
path2: Destination S3 path (s3://bucket/key).
**kwargs: Additional S3 copy parameters, as for :meth:`_cp_file`.
Returns:
False if ``path1`` is a directory, which is skipped; True if the
object was copied.
Raises:
ValueError: If trying to copy to a versioned file or copy buckets.
"""
# Parameters of the multipart copy, not of the S3 requests.
block_size = kwargs.pop("block_size", None)
max_workers = kwargs.pop("max_workers", None)
source = S3Path.parse(path1)
destination = S3Path.parse(path2)
if destination.version_id:
raise ValueError("Cannot copy to a versioned file.")
if not source.key or not destination.key:
raise ValueError("Cannot copy buckets.")
info1 = await self._info(path1)
if info1.get("type") == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY:
# Recursive copy() passes the directories too; see
# S3FileSystem.cp_file.
return False
size1 = info1.get("size", 0)
try:
if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE:
await asyncio.to_thread(self.core.copy_object, source, destination, **kwargs)
else:
await self._copy_object_with_multipart_upload(
source,
destination,
max_workers=max_workers,
block_size=block_size,
**kwargs,
)
finally:
# See S3FileSystem._copy_file.
self._sync_fs.invalidate_cache(path2)
return True
async def _copy_object_with_multipart_upload(
self,
source: S3Path,
destination: S3Path,
max_workers: int | None = None,
block_size: int | None = None,
**kwargs,
) -> None:
"""Copy an object with a multipart upload of its byte ranges.
See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part
and annotation copies run in parallel as asyncio tasks with
``asyncio.to_thread``. On a cancellation, the creation of the upload,
the running part copies and the completion are waited for, the
upload is aborted if it was created and has not completed, and the
cancellation is re-raised. A repeated cancellation returns without
stopping this cleanup.
Args:
source: Source S3 path, with the version ID to copy, if any.
destination: Destination S3 path.
max_workers: Maximum number of parallel requests.
block_size: Size in bytes of the copied ranges.
**kwargs: The CopyObject parameters of the copy; each request
receives those that it accepts.
Raises:
ValueError: If ``block_size`` is out of the part size limits, a
directive has an invalid value, or HeadObject reports no
size.
"""
max_workers = max_workers if max_workers else self._sync_fs.max_workers
plan = await asyncio.to_thread(
self.core.plan_multipart_copy, source, destination, block_size, **kwargs
)
if plan.fits_single_request:
# See S3FileSystem._copy_object_with_multipart_upload.
await asyncio.to_thread(self.core.copy_object, plan.source, plan.destination, **kwargs)
return
multipart_upload: S3MultipartUpload
semaphore = asyncio.Semaphore(max_workers)
failed = False
async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart | None:
nonlocal failed
async with semaphore:
if failed:
# The upload is being aborted; do not start more parts.
return None
try:
return await asyncio.to_thread(
self.core.upload_part_copy,
upload=multipart_upload,
part_number=i + 1,
source=plan.source,
range_=range_,
**plan.part_params,
)
except Exception:
# Set before the semaphore lets a waiting part start.
failed = True
raise
tasks: list[asyncio.Future[S3MultipartUploadPart | None]] = []
completion: asyncio.Task[S3CompleteMultipartUpload] | None = None
async def _abort() -> None:
await asyncio.wait([creation])
if creation.cancelled() or creation.exception() is not None:
# No upload was created.
return
# A part that is still copying when the upload is aborted may be
# stored after the abort, so wait for the running parts first.
await asyncio.gather(*tasks, return_exceptions=True)
if completion is not None:
await asyncio.wait([completion])
if not completion.cancelled() and completion.exception() is None:
# The upload completed despite the cancellation, so
# there is nothing to abort.
return
await asyncio.to_thread(
self._sync_fs._abort_multipart_upload,
creation.result(),
plan.abort_params,
)
# A task, so that _abort() can wait for an upload that is created
# after a cancellation; scheduled right before the try, so that
# nothing can fail between the two.
creation = asyncio.ensure_future(
asyncio.to_thread(
self.core.create_multipart_upload, plan.destination, **plan.create_params
)
)
try:
# shield keeps a cancellation from cancelling the creation, whose
# thread would keep running, so that _abort() can wait for it.
multipart_upload = await asyncio.shield(creation)
tasks = [asyncio.ensure_future(_upload_part(i, r)) for i, r in enumerate(plan.ranges)]
# Unlike gather, wait does not cancel the parts when this task is
# cancelled; their threads would keep copying, so they are waited
# for in _abort().
done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
for task in done:
if (error := task.exception()) is not None:
raise error
# The tasks are in part-number order.
parts = [task.result() for task in tasks]
completion = asyncio.ensure_future(
asyncio.to_thread(
self.core.complete_multipart_upload,
multipart_upload,
cast(list[S3MultipartUploadPart], parts),
**plan.complete_params,
)
)
# shield keeps a cancellation from cancelling the completion, whose
# thread would keep running, so that _abort() can wait for it.
completed = await asyncio.shield(completion)
except BaseException:
# Also on cancellation, as S3FileSystem._finish_multipart_upload
# does on an interrupt.
failed = True
cleanup = asyncio.ensure_future(_abort())
# A repeated cancellation of this task returns without stopping
# the cleanup; the event loop keeps only weak references to tasks.
_cleanup_tasks.add(cleanup)
cleanup.add_done_callback(_cleanup_tasks.discard)
await asyncio.shield(cleanup)
raise
failed = False
async def _copy_annotation(name: str) -> None:
nonlocal failed
async with semaphore:
if failed:
# Do not start more copies after one failed, as
# S3FileSystem does.
return
try:
await asyncio.to_thread(
self.core.copy_object_annotation,
name,
plan.source,
plan.destination,
completed.version_id,
completed.etag,
**kwargs,
)
except Exception:
failed = True
raise
results = await asyncio.gather(
*[_copy_annotation(name) for name in plan.annotations], return_exceptions=True
)
for result in results:
if isinstance(result, BaseException):
raise result
async def _find(
self,
path: str,
maxdepth: int | None = None,
withdirs: bool = False,
**kwargs,
) -> dict[str, S3Object] | list[str]:
detail = kwargs.pop("detail", False)
files = await asyncio.to_thread(
self._sync_fs._find, path, maxdepth=maxdepth, withdirs=withdirs, **kwargs
)
if detail:
return {f.name: f for f in files}
return [f.name for f in files]
async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> list[str]:
"""Expand glob patterns and directories into the paths they match.
See :meth:`S3FileSystem.expand_path`. The other paths are expanded by
fsspec's ``_expand_path()``, whose glob lists only the keys under the
literal part of the pattern.
Args:
path: S3 path, glob pattern, or list of them.
recursive: Whether to include the paths below the directories.
maxdepth: Maximum depth of the expansion, at least 1.
**kwargs: Additional arguments passed to fsspec's
``_expand_path()``, such as ``assume_literal``.
Returns:
The sorted matching paths.
Raises:
ValueError: If ``maxdepth`` is less than 1.
FileNotFoundError: If nothing matches.
"""
if maxdepth is not None and maxdepth < 1:
raise ValueError("maxdepth must be at least 1")
versions, others = self._sync_fs._split_version_paths(path)
out = {p for p in versions if not recursive or await self._exists(p)}
if others:
try:
out.update(
await super()._expand_path(
others, recursive=recursive, maxdepth=maxdepth, **kwargs
)
)
except FileNotFoundError:
if not out:
raise
if not out:
raise FileNotFoundError(path)
return sorted(out)
def _create_executor(self, max_workers: int) -> S3Executor:
"""Create the executor for the parallel operations of a file.
An instance created with ``asynchronous=True`` has no event loop of
its own, so its files run the operations in a thread pool.
Args:
max_workers: The maximum number of operations that run at once.
Returns:
An ``S3AioExecutor`` on the event loop of this filesystem, or an
``S3ThreadPoolExecutor`` if it has none.
"""
if self._loop is None:
return S3ThreadPoolExecutor(max_workers=max_workers)
return S3AioExecutor(loop=self._loop, max_workers=max_workers)
def _open(
self,
path: str,
mode: str = "rb",
block_size: int | None = None,
cache_type: str | None = None,
autocommit: bool = True,
cache_options: dict[Any, Any] | None = None,
**kwargs,
) -> AioS3File:
if block_size is None:
block_size = self._sync_fs.default_block_size
if cache_type is None:
cache_type = self._sync_fs.default_cache_type
max_workers = kwargs.pop("max_workers", self._sync_fs.max_workers)
# The parameters of the call take precedence over those of the
# filesystem; the caller's dictionary is not modified.
s3_additional_kwargs = {
**self._sync_fs.s3_additional_kwargs,
**kwargs.pop("s3_additional_kwargs", {}),
}
return AioS3File(
self._sync_fs,
path,
mode,
max_workers=max_workers,
executor=self._create_executor(max_workers=max_workers),
block_size=block_size,
cache_type=cache_type,
autocommit=autocommit,
cache_options=cache_options,
s3_additional_kwargs=s3_additional_kwargs,
**kwargs,
)
async def _rmdir(self, path: str) -> None:
await asyncio.to_thread(self._sync_fs.rmdir, path)
[docs]
def rmdir(self, path: str) -> None:
"""Remove an S3 bucket, which must be empty.
See :meth:`S3FileSystem.rmdir`.
Args:
path: S3 bucket path (e.g., "s3://bucket").
"""
self._sync_fs.rmdir(path)
[docs]
def sign(self, path: str, expiration: int = 3600, **kwargs) -> str:
"""Generate a presigned URL for S3 object access.
See :meth:`S3FileSystem.sign`.
Args:
path: S3 path (s3://bucket/key) to generate the URL for.
expiration: URL expiration time in seconds.
**kwargs: Additional parameters passed to :meth:`S3FileSystem.sign`.
Returns:
The presigned URL.
"""
return cast(str, self._sync_fs.sign(path, expiration=expiration, **kwargs))
[docs]
def getxattr(self, path: str, attr_name: str, **kwargs) -> str | None:
"""Get an attribute from the user-defined metadata of the path.
See :meth:`S3FileSystem.getxattr`.
Args:
path: S3 path (s3://bucket/key) to get the attribute for.
attr_name: The name of the attribute.
**kwargs: Additional parameters passed to the HeadObject API.
Returns:
The value of the attribute, or None if the attribute is not set.
"""
return self._sync_fs.getxattr(path, attr_name, **kwargs)
[docs]
def setxattr(self, path: str, copy_kwargs: dict[str, Any] | None = None, **kwargs) -> None:
"""Set the user-defined metadata of the path.
See :meth:`S3FileSystem.setxattr`.
Args:
path: S3 path (s3://bucket/key) to set metadata for.
copy_kwargs: Additional parameters to use for the underlying
CopyObject API call.
**kwargs: Key-value pairs of metadata to set; a None value deletes the
key.
"""
self._sync_fs.setxattr(path, copy_kwargs=copy_kwargs, **kwargs)
[docs]
def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None:
"""Set the canned ACL of a bucket or key.
See :meth:`S3FileSystem.chmod`.
Args:
path: S3 path (s3://bucket or s3://bucket/key) to set the ACL on.
acl: The canned ACL to apply.
recursive: Whether to apply the ACL to all keys below the path too.
**kwargs: Additional parameters passed to the PutObjectAcl or
PutBucketAcl API.
"""
self._sync_fs.chmod(path, acl, recursive=recursive, **kwargs)
[docs]
def object_version_info(
self, path: str, delete_markers: bool = False, **kwargs
) -> list[S3ObjectVersion]:
"""List the versions of the object or of the objects under the path.
See :meth:`S3FileSystem.object_version_info`.
Args:
path: S3 path (s3://bucket/key or a key prefix) to list the versions for.
delete_markers: Whether to include delete markers in the result.
**kwargs: Additional parameters passed to the ListObjectVersions API.
Returns:
List of S3ObjectVersion instances describing the versions.
"""
return self._sync_fs.object_version_info(path, delete_markers=delete_markers, **kwargs)
[docs]
def list_multipart_uploads(self, path: str) -> list[S3MultipartUpload]:
"""List in-progress (incomplete) multipart uploads in a bucket.
See :meth:`S3FileSystem.list_multipart_uploads`.
Args:
path: S3 bucket or key path (e.g., "s3://bucket" or "s3://bucket/prefix").
Returns:
List of S3MultipartUpload instances describing the uploads.
"""
return self._sync_fs.list_multipart_uploads(path)
[docs]
def clear_multipart_uploads(self, path: str) -> None:
"""Abort any incomplete multipart uploads in the bucket.
See :meth:`S3FileSystem.clear_multipart_uploads`.
Args:
path: S3 bucket or key path (e.g., "s3://bucket" or "s3://bucket/prefix").
"""
self._sync_fs.clear_multipart_uploads(path)
[docs]
def checksum(self, path: str, **kwargs) -> int:
"""Get the checksum of an S3 object or directory.
See :meth:`S3FileSystem.checksum`.
Args:
path: S3 path (s3://bucket/key) to get the checksum for.
**kwargs: Additional arguments passed to :meth:`S3FileSystem.checksum`.
Returns:
Integer checksum derived from the ETag or the directory token.
"""
return cast(int, self._sync_fs.checksum(path, **kwargs))
[docs]
def created(self, path: str) -> datetime:
"""Return the creation time of the path.
See :meth:`S3FileSystem.created`.
Args:
path: S3 path (s3://bucket/key).
Returns:
The last-modified time of the object.
"""
return self._sync_fs.created(path)
[docs]
def modified(self, path: str) -> datetime:
"""Return the last-modified time of the path.
See :meth:`S3FileSystem.modified`.
Args:
path: S3 path (s3://bucket/key).
Returns:
The last-modified time of the object.
"""
return self._sync_fs.modified(path)
[docs]
def invalidate_cache(self, path: str | None = None) -> None:
"""Remove the cached entries of the path and its parent paths.
See :meth:`S3FileSystem.invalidate_cache`.
Args:
path: The path to invalidate. If None, clear the whole cache.
"""
self._sync_fs.invalidate_cache(path)
async def _touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]:
return await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs)
[docs]
def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]:
"""Create an empty object with PutObject.
See :meth:`S3FileSystem.touch`.
Args:
path: S3 path (s3://bucket/key) of the object.
truncate: If True, replace an existing object with an empty one;
if False, raise if the object exists.
**kwargs: Additional parameters passed to the PutObject API.
Returns:
The PutObject response as a dictionary.
"""
return self._sync_fs.touch(path, truncate=truncate, **kwargs)
[docs]
class AioS3File(S3File):
"""Async-aware S3 file handle using ``S3AioExecutor``.
Functionally identical to ``S3File``; exists as a distinct type for
``isinstance`` checks and to document the async execution model.
All parallel operations (range reads, multipart uploads) are dispatched
through the ``S3Executor`` interface — the ``S3AioExecutor``
provided by ``AioS3FileSystem`` dispatches them through the event loop with
``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file.
For an ``AioS3FileSystem`` created with ``asynchronous=True``, it is an
``S3ThreadPoolExecutor``.
"""