megfile 4.2.3__py3-none-any.whl → 4.2.5__py3-none-any.whl

megfile/__init__.py

@@ -206,6 +206,16 @@ from megfile.stdio import is_stdio, stdio_open
 from megfile.stdio_path import StdioPath
 from megfile.version import VERSION as __version__  # noqa: F401
 
+try:
+    from megfile.sftp2_path import Sftp2Path
+except ImportError:
+    Sftp2Path = None
+
+try:
+    from megfile.webdav_path import WebdavPath
+except ImportError:
+    WebdavPath = None
+
 __all__ = [
     "smart_access",
     "smart_cache",

megfile/cli.py

@@ -47,7 +47,7 @@ from megfile.smart import (
     smart_unlink,
 )
 from megfile.smart_path import SmartPath
-from megfile.utils import get_human_size
+from megfile.utils import copyfileobj_multi, get_human_size
 from megfile.version import VERSION
 
 options = {}
@@ -646,14 +646,10 @@ def to(path: str, append: bool, stdout: bool):
         smart_open(path, mode) as f,
         smart_open("stdio://1", "wb") as stdout_fd,
     ):
-        length = 16 * 1024
-        while True:
-            buf = stdin.read(length)
-            if not buf:
-                break
-            f.write(buf)
-            if stdout:
-                stdout_fd.write(buf)
+        destinations = [f]
+        if stdout:
+            destinations.append(stdout_fd)
+        copyfileobj_multi(stdin, destinations)
 
 
 @cli.command(short_help="Produce an md5sum file for all the objects in the path.")

megfile/config.py

@@ -83,6 +83,7 @@ if READER_BLOCK_SIZE <= 0:
 READER_MAX_BUFFER_SIZE = parse_quantity(
     os.getenv("MEGFILE_READER_MAX_BUFFER_SIZE") or 128 * 2**20
 )
+READER_LAZY_PREFETCH = parse_boolean(os.getenv("MEGFILE_READER_LAZY_PREFETCH"), False)
 
 # Multi-upload in aws s3 has a maximum of 10,000 parts,
 # so the maximum supported file size is MEGFILE_WRITE_BLOCK_SIZE * 10,000,
@@ -105,6 +106,10 @@ GLOBAL_MAX_WORKERS = int(os.getenv("MEGFILE_MAX_WORKERS") or 8)
 
 NEWLINE = ord("\n")
 
+# Default buffer sizes for various operations
+DEFAULT_COPY_BUFFER_SIZE = 16 * 1024  # 16KB, same as shutil.copyfileobj
+DEFAULT_HASH_BUFFER_SIZE = 4 * 1024  # 4KB for hash calculations
+
 S3_CLIENT_CACHE_MODE = os.getenv("MEGFILE_S3_CLIENT_CACHE_MODE") or "thread_local"
 
 DEFAULT_MAX_RETRY_TIMES = int(os.getenv("MEGFILE_MAX_RETRY_TIMES") or 10)

megfile/fs.py

@@ -1,7 +1,7 @@
 import os
 from stat import S_ISDIR as stat_isdir
 from stat import S_ISLNK as stat_islnk
-from typing import BinaryIO, Callable, Iterator, List, Optional, Tuple
+from typing import IO, BinaryIO, Callable, Iterator, List, Optional, Tuple
 
 from megfile.fs_path import (
     FSPath,
@@ -52,6 +52,7 @@ __all__ = [
     "fs_islink",
     "fs_ismount",
     "fs_save_as",
+    "fs_open",
 ]
 
 
@@ -612,3 +613,15 @@ def fs_move(src_path: PathLike, dst_path: PathLike, overwrite: bool = True) -> N
     :param overwrite: whether or not overwrite file when exists
     """
     return fs_rename(src_path, dst_path, overwrite)
+
+
+def fs_open(path: PathLike, mode: str = "r", **kwargs) -> IO:
+    """
+    Open file on fs
+
+    :param path: Given path
+    :param mode: File open mode, like built-in open function
+    :param buffering: Buffering policy, like built-in open function
+    :returns: A file-like object
+    """
+    return FSPath(path).open(mode, **kwargs)

megfile/fs_path.py

@@ -17,6 +17,7 @@ from megfile.interfaces import (
     Access,
     ContextIterator,
     FileEntry,
+    FileLike,
     PathLike,
     StatResult,
     URIPath,
@@ -27,7 +28,7 @@ from megfile.lib.glob import iglob
 from megfile.lib.joinpath import path_join
 from megfile.lib.url import get_url_scheme
 from megfile.smart_path import SmartPath
-from megfile.utils import calculate_md5
+from megfile.utils import calculate_md5, copyfd
 
 __all__ = [
     "FSPath",
@@ -85,6 +86,36 @@ def _fs_rename_file(
     shutil.move(src_path, dst_path)
 
 
+class WrapAtomic(FileLike):
+    __atomic__ = True
+
+    def __init__(self, fileobj):
+        self.fileobj = fileobj
+        self.temp_name = f"{self.name}.temp"
+        os.rename(self.name, self.temp_name)
+
+    @property
+    def name(self):
+        return self.fileobj.name
+
+    @property
+    def mode(self):
+        return self.fileobj.mode
+
+    def _close(self):
+        self.fileobj.close()
+        os.rename(self.temp_name, self.name)
+
+    def _abort(self):
+        try:
+            os.unlink(self.temp_name)
+        except FileNotFoundError:
+            pass
+
+    def __getattr__(self, name: str):
+        return getattr(self.fileobj, name)
+
+
 @SmartPath.register
 class FSPath(URIPath):
     """file protocol
@@ -627,9 +658,11 @@ class FSPath(URIPath):
         """
         self._check_int_path()
 
-        if missing_ok and not self.exists():
-            return
-        os.unlink(self.path_without_protocol)  # pyre-ignore[6]
+        try:
+            os.unlink(self.path_without_protocol)  # pyre-ignore[6]
+        except FileNotFoundError:
+            if not missing_ok:
+                raise
 
     def walk(
         self, followlinks: bool = False
@@ -737,15 +770,7 @@ class FSPath(URIPath):
     ):
         if isinstance(self.path_without_protocol, int):
             with open(fspath(dst_path), "wb") as fdst:
-                # This magic number is copied from  copyfileobj
-                length = 16 * 1024
-                while True:
-                    buf = os.read(self.path_without_protocol, length)  # pyre-ignore[6]
-                    if not buf:
-                        break
-                    fdst.write(buf)
-                    if callback:
-                        callback(len(buf))
+                copyfd(self.path_without_protocol, fdst, callback)
         else:
             shutil.copy2(
                 self.path_without_protocol,  # pyre-ignore[6]
@@ -925,11 +950,12 @@ class FSPath(URIPath):
     def open(
         self,
         mode: str = "r",
-        buffering=-1,
-        encoding=None,
-        errors=None,
-        newline=None,
-        closefd=True,
+        buffering: int = -1,
+        encoding: Optional[str] = None,
+        errors: Optional[str] = None,
+        newline: Optional[str] = None,
+        closefd: bool = True,
+        atomic: bool = False,
         **kwargs,
     ) -> IO:
         if not isinstance(self.path_without_protocol, int) and (
@@ -940,7 +966,7 @@ class FSPath(URIPath):
                     self.path_without_protocol  # pyre-ignore[6]
                 )
             ).mkdir(parents=True, exist_ok=True)
-        return io.open(
+        fp = io.open(
             self.path_without_protocol,
             mode,
             buffering=buffering,
@@ -949,6 +975,9 @@ class FSPath(URIPath):
             newline=newline,
             closefd=closefd,
         )
+        if atomic and ("w" in mode or "x" in mode or "a" in mode):
+            return WrapAtomic(fp)
+        return fp
 
     @cached_property
     def parts(self) -> Tuple[str, ...]:

megfile/interfaces.py

@@ -1,6 +1,7 @@
 import os
 from abc import ABC, abstractmethod
 from io import IOBase, UnsupportedOperation
+from logging import getLogger as get_logger
 from typing import IO, AnyStr, Iterable, List, Optional
 
 from megfile.pathlike import (
@@ -31,6 +32,8 @@ __all__ = [
     "URIPath",
 ]
 
+_logger = get_logger(__name__)
+
 
 def fullname(o):
     klass = o.__class__
@@ -43,16 +46,28 @@ def fullname(o):
 # 1. Default value of closed is False
 # 2. closed is set to True when close() are called
 # 3. close() will only be called once
+# 4. atomic means the file-like object should not be closed automatically
+#    when an exception is raised in the context manager or when the object is
+#    garbage collected.
+# 5. atomic is False by default
 class Closable(ABC):
     @property
     def closed(self) -> bool:
         """Return True if the file-like object is closed."""
         return getattr(self, "__closed__", False)
 
+    @property
+    def atomic(self) -> bool:
+        """Return True if the file-like object is atomic."""
+        return getattr(self, "__atomic__", False)
+
     @abstractmethod
     def _close(self) -> None:
         pass  # pragma: no cover
 
+    def _abort(self) -> None:
+        pass
+
     def close(self) -> None:
         """Flush and close the file-like object.
 
@@ -66,6 +81,24 @@ class Closable(ABC):
         return self
 
     def __exit__(self, type, value, traceback) -> None:
+        if self.atomic and value is not None:
+            from megfile.errors import full_error_message
+
+            _logger.warning(
+                f"skip closing atomic file-like object: {self}, "
+                f"since error encountered: {full_error_message(value)}"
+            )
+            self._abort()
+            return
+        self.close()
+
+    def __del__(self):
+        if self.atomic:
+            _logger.warning(
+                f"skip closing atomic file-like object before deletion: {self}"
+            )
+            self._abort()
+            return
         self.close()
 
 

megfile/lib/base_prefetch_reader.py

@@ -82,9 +82,9 @@ class BasePrefetchReader(Readable[bytes], Seekable, ABC):
 
         self._offset = 0
         self._cached_buffer = None
-        self._block_index = None  # Current block index
+        self._block_index = 0  # Current block index
+        self._cached_offset = 0  # Current offset in the current block
         self._seek_history = []
-
         self._seek_buffer(0)
 
         _logger.debug("open file: %r, mode: %s" % (self.name, self.mode))
@@ -98,7 +98,9 @@ class BasePrefetchReader(Readable[bytes], Seekable, ABC):
         return self._process_local("futures", self._get_futures)
 
     def _get_futures(self):
-        return LRUCacheFutureManager()
+        futures = LRUCacheFutureManager()
+        futures.register(self.name)
+        return futures
 
     @property
     @abstractmethod
@@ -207,9 +209,8 @@ class BasePrefetchReader(Readable[bytes], Seekable, ABC):
         if size == 0 or self._offset >= self._content_size:
             return b""
 
-        data = self._fetch_response(start=self._offset, end=self._offset + size - 1)[
-            "Body"
-        ].read()
+        resp = self._fetch_response(start=self._offset, end=self._offset + size - 1)
+        data = resp["Body"].read()
         self.seek(size, os.SEEK_CUR)
         return data
 
@@ -369,12 +370,17 @@ class BasePrefetchReader(Readable[bytes], Seekable, ABC):
 class LRUCacheFutureManager(OrderedDict):
     def __init__(self):
         super().__init__()
+        self._name = None
+
+    def register(self, name):
+        self._name = name
 
     def submit(self, executor, key, *args, **kwargs):
         if key in self:
             self.move_to_end(key, last=True)
             return
         self[key] = executor.submit(*args, **kwargs)
+        _logger.debug("submit future: %r, key: %r" % (self._name, key))
 
     @property
     def finished(self):
@@ -385,7 +391,12 @@ class LRUCacheFutureManager(OrderedDict):
         return self[key].result()
 
     def cleanup(self, block_capacity: int):
+        keys = []
         while len(self) > block_capacity:
-            _, future = self.popitem(last=False)
+            key, future = self.popitem(last=False)
+            keys.append(key)
             if not future.done():
                 future.cancel()
+        if keys:
+            _logger.debug("cleanup futures: %r, keys: %s" % (self._name, keys))
+        return keys

megfile/lib/joinpath.py

@@ -33,3 +33,16 @@ def uri_join(path: str, *other_paths: str) -> str:
 
     # Imp. 3
     # return '/'.join((path, *other_paths))
+
+
+def uri_norm(path: str) -> str:
+    parts = path.split("/")
+    new_parts = []
+    for part in parts:
+        if part == ".":
+            continue
+        if part == ".." and new_parts and new_parts[-1] != "..":
+            new_parts.pop()
+        else:
+            new_parts.append(part)
+    return "/".join(new_parts)

megfile/lib/s3_buffered_writer.py

@@ -53,11 +53,13 @@ class S3BufferedWriter(Writable[bytes]):
         max_buffer_size: int = WRITER_MAX_BUFFER_SIZE,
         max_workers: Optional[int] = None,
         profile_name: Optional[str] = None,
+        atomic: bool = False,
     ):
         self._bucket = bucket
         self._key = key
         self._client = s3_client
         self._profile_name = profile_name
+        self.__atomic__ = atomic
 
         # user maybe put block_size with 'numpy.uint64' type
         self._base_block_size = int(block_size)
@@ -213,6 +215,17 @@ class S3BufferedWriter(Writable[bytes]):
         if not self._is_global_executor:
             self._executor.shutdown()
 
+    def _abort(self):
+        _logger.debug("abort file: %r" % self.name)
+
+        if self._is_multipart:
+            with raise_s3_error(self.name):
+                self._client.abort_multipart_upload(
+                    Bucket=self._bucket, Key=self._key, UploadId=self._upload_id
+                )
+
+        self._shutdown()
+
     def _close(self):
         _logger.debug("close file: %r" % self.name)
 

megfile/lib/s3_limited_seekable_writer.py

@@ -33,6 +33,7 @@ class S3LimitedSeekableWriter(S3BufferedWriter, Seekable):
         max_buffer_size: int = WRITER_MAX_BUFFER_SIZE,
         max_workers: Optional[int] = None,
         profile_name: Optional[str] = None,
+        atomic: bool = False,
     ):
         super().__init__(
             bucket,
@@ -42,6 +43,7 @@ class S3LimitedSeekableWriter(S3BufferedWriter, Seekable):
             max_buffer_size=max_buffer_size,
             max_workers=max_workers,
             profile_name=profile_name,
+            atomic=atomic,
         )
 
         self._head_block_size = head_block_size or block_size

megfile/lib/s3_prefetch_reader.py

@@ -4,6 +4,7 @@ from typing import Optional
 
 from megfile.config import (
     READER_BLOCK_SIZE,
+    READER_LAZY_PREFETCH,
     READER_MAX_BUFFER_SIZE,
     S3_MAX_RETRY_TIMES,
 )
@@ -62,7 +63,7 @@ class S3PrefetchReader(BasePrefetchReader):
         )
 
     def _get_content_size(self):
-        if self._block_capacity <= 0:
+        if self._block_capacity <= 0 or READER_LAZY_PREFETCH:
             response = self._client.head_object(Bucket=self._bucket, Key=self._key)
             self._content_etag = response.get("ETag")
             return int(response["ContentLength"])

megfile/lib/s3_share_cache_reader.py

@@ -101,16 +101,21 @@ class ShareCacheFutureManager(LRUCacheFutureManager):
         super().__init__()
         self._references = Counter()
 
-    def register(self, key):
-        self._references[key] += 1
-
-    def unregister(self, key):
-        self._references[key] -= 1
-        if self._references[key] == 0:
-            self._references.pop(key)
-            for key_tuple in list(self):
-                if key_tuple[0] != key:
+    def register(self, name):
+        self._references[name] += 1
+        _logger.debug("register reader: %r, count: %d" % (name, self._references[name]))
+
+    def unregister(self, name):
+        self._references[name] -= 1
+        _logger.debug(
+            "unregister reader: %r, count: %d" % (name, self._references[name])
+        )
+        if self._references[name] == 0:
+            self._references.pop(name)
+            for key in list(self):
+                if key[0] != name:
                     continue
-                future = self.pop(key_tuple)
+                future = self.pop(key)
                 if not future.done():
                     future.cancel()  # pragma: no cover
+            _logger.debug("cleanup all futures of reader: %r" % name)

megfile/s3_path.py

@@ -230,7 +230,7 @@ def get_endpoint_url(profile_name: Optional[str] = None) -> str:
     config_endpoint_url = config.get("s3", {}).get("endpoint_url")
     config_endpoint_url = config_endpoint_url or config.get("endpoint_url")
     if config_endpoint_url:
-        warning_endpoint_url("~/.aws/config", config_endpoint_url)
+        warning_endpoint_url("~/.aws/config or ~/.aws/credentials", config_endpoint_url)
         return config_endpoint_url
     return endpoint_url
 
@@ -937,6 +937,7 @@ def s3_buffered_open(
     buffered: bool = False,
     share_cache_key: Optional[str] = None,
     cache_path: Optional[str] = None,
+    atomic: bool = False,
 ) -> IO:
     """Open an asynchronous prefetch reader, to support fast sequential read
 
@@ -969,6 +970,8 @@ def s3_buffered_open(
         (both file head part and tail part can seek block_size).
         Notes: This parameter are valid only for write-handle.
         Read-handle support arbitrary seek
+    :param buffered: If you are operating pickle file without .pkl or .pickle extension,
+        please set this to True to avoid the performance issue.
     :returns: An opened File object
     :raises: S3FileNotFoundError
     """
@@ -1000,6 +1003,7 @@ def s3_buffered_open(
         )
 
     if mode == "rb":
+        block_size = block_size or READER_BLOCK_SIZE
         if share_cache_key is not None:
             reader = S3ShareCacheReader(
                 bucket,
@@ -1008,7 +1012,7 @@ def s3_buffered_open(
                 s3_client=client,
                 max_retries=max_retries,
                 max_workers=max_workers,
-                block_size=block_size or READER_BLOCK_SIZE,
+                block_size=block_size,
                 block_forward=block_forward,
                 profile_name=s3_url._profile_name,
             )
@@ -1023,13 +1027,14 @@ def s3_buffered_open(
                 max_workers=max_workers,
                 max_buffer_size=max_buffer_size,
                 block_forward=block_forward,
-                block_size=block_size or READER_BLOCK_SIZE,
+                block_size=block_size,
                 profile_name=s3_url._profile_name,
             )
         if buffered or _is_pickle(reader):
             reader = io.BufferedReader(reader)  # type: ignore
         return reader
 
+    block_size = block_size or WRITER_BLOCK_SIZE
     if limited_seekable:
         if max_buffer_size is None:
             max_buffer_size = WRITER_MAX_BUFFER_SIZE
@@ -1038,9 +1043,10 @@ def s3_buffered_open(
             key,
             s3_client=client,
             max_workers=max_workers,
-            block_size=block_size or WRITER_BLOCK_SIZE,
+            block_size=block_size,
             max_buffer_size=max_buffer_size,
             profile_name=s3_url._profile_name,
+            atomic=atomic,
         )
     else:
         if max_buffer_size is None:
@@ -1050,9 +1056,10 @@ def s3_buffered_open(
             key,
             s3_client=client,
             max_workers=max_workers,
-            block_size=block_size or WRITER_BLOCK_SIZE,
+            block_size=block_size,
             max_buffer_size=max_buffer_size,
             profile_name=s3_url._profile_name,
+            atomic=atomic,
         )
     if buffered or _is_pickle(writer):
         writer = io.BufferedWriter(writer)  # type: ignore