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

megfile/sftp_path.py

@@ -1,4 +1,5 @@
 import atexit
+import base64
 import hashlib
 import io
 import os
@@ -29,6 +30,11 @@ _logger = get_logger(__name__)
 __all__ = [
     "SftpPath",
     "is_sftp",
+    "sftp_add_host_key",
+    "sftp_concat",
+    "sftp_copy",
+    "sftp_download",
+    "sftp_upload",
 ]
 
 SFTP_USERNAME = "SFTP_USERNAME"
@@ -64,7 +70,8 @@ def get_private_key():
         private_key_path = os.getenv(SFTP_PRIVATE_KEY_PATH)
         if not os.path.exists(private_key_path):
             raise FileNotFoundError(
-                f"Private key file not exist: '{SFTP_PRIVATE_KEY_PATH}'"
+                "Private key file not exist, "
+                f"path:{private_key_path}, env:'{SFTP_PRIVATE_KEY_PATH}'"
             )
         return key_with_types[key_type].from_private_key_file(
             private_key_path, password=os.getenv(SFTP_PRIVATE_KEY_PASSWORD)
@@ -335,6 +342,236 @@ def _sftp_scan_pairs(
         yield src_file_path, dst_file_path
 
 
+def _check_input(input_str: str, fingerprint: str, times: int = 0) -> bool:
+    answers = input_str.strip()
+    if answers.lower() in ("yes", "y") or answers == fingerprint:
+        return True
+    elif answers.lower() in ("no", "n"):
+        return False
+    elif times >= 10:
+        _logger.warning("Retried more than 10 times, give up")
+        return False
+    else:
+        input_str = input("Please type 'yes', 'no' or the fingerprint: ")
+        return _check_input(input_str, fingerprint, times=times + 1)
+
+
+def _prompt_add_to_known_hosts(hostname, key) -> bool:
+    fingerprint = hashlib.sha256(key.asbytes()).digest()
+    fingerprint = f"SHA256:{base64.b64encode(fingerprint).decode('utf-8')}"
+    answers = input(f"""The authenticity of host '{hostname}' can't be established.
+{key.get_name().upper()} key fingerprint is {fingerprint}.
+This key is not known by any other names.
+Are you sure you want to continue connecting (yes/no/[fingerprint])? """)
+    return _check_input(answers, fingerprint)
+
+
+def sftp_add_host_key(
+    hostname: str,
+    port: int = 22,
+    prompt: bool = False,
+    host_key_path: Optional["str"] = None,
+):
+    """Add a host key to known_hosts.
+
+    :param hostname: hostname
+    :param port: port, default is 22
+    :param prompt: If True, requires user input of 'yes' or 'no' to decide whether to
+        add this host key
+    :param host_key_path: path of known_hosts, default is ~/.ssh/known_hosts
+    """
+    if not host_key_path:
+        host_key_path = os.path.expanduser("~/.ssh/known_hosts")
+
+    if not os.path.exists(host_key_path):
+        dirname = os.path.dirname(host_key_path)
+        if dirname and dirname != ".":
+            os.makedirs(dirname, exist_ok=True, mode=0o700)
+        with open(host_key_path, "w"):
+            pass
+        os.chmod(host_key_path, 0o600)
+
+    host_key = paramiko.hostkeys.HostKeys(host_key_path)
+    if host_key.lookup(hostname):
+        return
+
+    transport = paramiko.Transport(
+        (
+            hostname,
+            port,
+        )
+    )
+    transport.connect()
+    key = transport.get_remote_server_key()
+    transport.close()
+
+    if prompt:
+        result = _prompt_add_to_known_hosts(hostname, key)
+        if not result:
+            return
+
+    host_key.add(hostname, key.get_name(), key)
+    host_key.save(host_key_path)
+
+
+def sftp_concat(src_paths: List[PathLike], dst_path: PathLike) -> None:
+    """Concatenate sftp files to one file.
+
+    :param src_paths: Given source paths
+    :param dst_path: Given destination path
+    """
+    dst_path_obj = SftpPath(dst_path)
+
+    def get_real_path(path: PathLike) -> str:
+        return SftpPath(path)._real_path
+
+    command = ["cat", *map(get_real_path, src_paths), ">", get_real_path(dst_path)]
+    exec_result = dst_path_obj._exec_command(command)
+    if exec_result.returncode != 0:
+        _logger.error(exec_result.stderr)
+        raise OSError(f"Failed to concat {src_paths} to {dst_path}")
+
+
+def sftp_copy(
+    src_path: PathLike,
+    dst_path: PathLike,
+    callback: Optional[Callable[[int], None]] = None,
+    followlinks: bool = False,
+    overwrite: bool = True,
+):
+    """
+    Copy the file to the given destination path.
+
+    :param src_path: Given path
+    :param dst_path: The destination path to copy the file to.
+    :param callback: An optional callback function that takes an integer parameter
+        and is called periodically during the copy operation to report the number
+        of bytes copied.
+    :param followlinks: Whether to follow symbolic links when copying directories.
+    :raises IsADirectoryError: If the source is a directory.
+    :raises OSError: If there is an error copying the file.
+    """
+    return SftpPath(src_path).copy(dst_path, callback, followlinks, overwrite)
+
+
+def sftp_download(
+    src_url: PathLike,
+    dst_url: PathLike,
+    callback: Optional[Callable[[int], None]] = None,
+    followlinks: bool = False,
+    overwrite: bool = True,
+):
+    """
+    Downloads a file from sftp to local filesystem.
+
+    :param src_url: source sftp path
+    :param dst_url: target fs path
+    :param callback: Called periodically during copy, and the input parameter is
+        the data size (in bytes) of copy since the last call
+    :param followlinks: False if regard symlink as file, else True
+    :param overwrite: whether or not overwrite file when exists, default is True
+    """
+    from megfile.fs_path import FSPath, is_fs
+
+    if not is_fs(dst_url):
+        raise OSError(f"dst_url is not fs path: {dst_url}")
+    if not is_sftp(src_url) and not isinstance(src_url, SftpPath):
+        raise OSError(f"src_url is not sftp path: {src_url}")
+
+    dst_path = FSPath(dst_url)
+    if not overwrite and dst_path.exists():
+        return
+
+    if isinstance(src_url, SftpPath):
+        src_path: SftpPath = src_url
+    else:
+        src_path: SftpPath = SftpPath(src_url)
+
+    if followlinks and src_path.is_symlink():
+        src_path = src_path.readlink()
+    if src_path.is_dir():
+        raise IsADirectoryError("Is a directory: %r" % src_url)
+    if str(dst_url).endswith("/"):
+        raise IsADirectoryError("Is a directory: %r" % dst_url)
+
+    dst_path.parent.makedirs(exist_ok=True)
+
+    sftp_callback = None
+    if callback:
+        bytes_transferred_before = 0
+
+        def sftp_callback(bytes_transferred: int, _total_bytes: int):
+            nonlocal bytes_transferred_before
+            callback(bytes_transferred - bytes_transferred_before)  # pyre-ignore[29]
+            bytes_transferred_before = bytes_transferred
+
+    src_path._client.get(
+        src_path._real_path, dst_path.path_without_protocol, callback=sftp_callback
+    )
+
+    src_stat = src_path.stat()
+    dst_path.utime(src_stat.st_atime, src_stat.st_mtime)
+    dst_path.chmod(src_stat.st_mode)
+
+
+def sftp_upload(
+    src_url: PathLike,
+    dst_url: PathLike,
+    callback: Optional[Callable[[int], None]] = None,
+    followlinks: bool = False,
+    overwrite: bool = True,
+):
+    """
+    Uploads a file from local filesystem to sftp server.
+
+    :param src_url: source fs path
+    :param dst_url: target sftp path
+    :param callback: Called periodically during copy, and the input parameter is
+        the data size (in bytes) of copy since the last call
+    :param overwrite: whether or not overwrite file when exists, default is True
+    """
+    from megfile.fs_path import FSPath, is_fs
+
+    if not is_fs(src_url):
+        raise OSError(f"src_url is not fs path: {src_url}")
+    if not is_sftp(dst_url) and not isinstance(dst_url, SftpPath):
+        raise OSError(f"dst_url is not sftp path: {dst_url}")
+
+    if followlinks and os.path.islink(src_url):
+        src_url = os.readlink(src_url)
+    if os.path.isdir(src_url):
+        raise IsADirectoryError("Is a directory: %r" % src_url)
+    if str(dst_url).endswith("/"):
+        raise IsADirectoryError("Is a directory: %r" % dst_url)
+
+    src_path = FSPath(src_url)
+    if isinstance(dst_url, SftpPath):
+        dst_path: SftpPath = dst_url
+    else:
+        dst_path: SftpPath = SftpPath(dst_url)
+    if not overwrite and dst_path.exists():
+        return
+
+    dst_path.parent.makedirs(exist_ok=True)
+
+    sftp_callback = None
+    if callback:
+        bytes_transferred_before = 0
+
+        def sftp_callback(bytes_transferred: int, _total_bytes: int):
+            nonlocal bytes_transferred_before
+            callback(bytes_transferred - bytes_transferred_before)  # pyre-ignore[29]
+            bytes_transferred_before = bytes_transferred
+
+    dst_path._client.put(
+        src_path.path_without_protocol, dst_path._real_path, callback=sftp_callback
+    )
+
+    src_stat = src_path.stat()
+    dst_path.utime(src_stat.st_atime, src_stat.st_mtime)
+    dst_path.chmod(src_stat.st_mode)
+
+
 @SmartPath.register
 class SftpPath(URIPath):
     """sftp protocol

megfile/smart.py

@@ -16,14 +16,13 @@ from typing import (
 from tqdm import tqdm
 
 from megfile.errors import S3UnknownError
-from megfile.fs import (
+from megfile.fs_path import (
     fs_copy,
-    fs_listdir,
-    fs_scandir,
     is_fs,
 )
 from megfile.interfaces import (
     Access,
+    ContextIterator,
     FileCacher,
     FileEntry,
     NullCacher,
@@ -34,7 +33,7 @@ from megfile.lib.combine_reader import CombineReader
 from megfile.lib.compare import get_sync_type, is_same_file
 from megfile.lib.compat import fspath
 from megfile.lib.glob import globlize, ungloblize
-from megfile.s3 import (
+from megfile.s3_path import (
     is_s3,
     s3_concat,
     s3_copy,
@@ -43,7 +42,7 @@ from megfile.s3 import (
     s3_open,
     s3_upload,
 )
-from megfile.sftp import sftp_concat, sftp_copy, sftp_download, sftp_upload
+from megfile.sftp_path import sftp_concat, sftp_copy, sftp_download, sftp_upload
 from megfile.smart_path import SmartPath, get_traditional_path
 from megfile.utils import combine, copyfileobj, generate_cache_path
 
@@ -172,7 +171,7 @@ def smart_listdir(path: Optional[PathLike] = None) -> List[str]:
     :raises: FileNotFoundError, NotADirectoryError
     """
     if path is None:
-        return fs_listdir()
+        return sorted(os.listdir(path))
     return SmartPath(path).listdir()
 
 
@@ -185,7 +184,28 @@ def smart_scandir(path: Optional[PathLike] = None) -> Iterator[FileEntry]:
     :raises: FileNotFoundError, NotADirectoryError
     """
     if path is None:
-        return fs_scandir()
+
+        def create_generator():
+            from stat import S_ISDIR as stat_isdir
+            from stat import S_ISLNK as stat_islnk
+
+            with os.scandir(None) as entries:
+                for entry in entries:
+                    stat = entry.stat()
+                    yield FileEntry(
+                        entry.name,
+                        entry.path,
+                        StatResult(
+                            size=stat.st_size,
+                            ctime=stat.st_ctime,
+                            mtime=stat.st_mtime,
+                            isdir=stat_isdir(stat.st_mode),
+                            islnk=stat_islnk(stat.st_mode),
+                            extra=stat,
+                        ),
+                    )
+
+        return ContextIterator(create_generator())
     return SmartPath(path).scandir()
 
 
@@ -359,8 +379,8 @@ def smart_copy(
     if smart_islink(src_path) and is_s3(dst_path) and not followlinks:
         return
 
-    src_protocol, _ = SmartPath._extract_protocol(src_path)
-    dst_protocol, _ = SmartPath._extract_protocol(dst_path)
+    src_protocol = SmartPath._extract_protocol(src_path)
+    dst_protocol = SmartPath._extract_protocol(dst_path)
 
     copy_func = _get_copy_func(src_protocol, dst_protocol)
 
@@ -406,8 +426,8 @@ def _smart_sync_single_file(items: dict):
         # this function is equal to smart_copy
         dst_abs_file_path = dst_root_path
 
-    src_protocol, _ = SmartPath._extract_protocol(src_file_path)
-    dst_protocol, _ = SmartPath._extract_protocol(dst_abs_file_path)
+    src_protocol = SmartPath._extract_protocol(src_file_path)
+    dst_protocol = SmartPath._extract_protocol(dst_abs_file_path)
     should_sync = True
     try:
         if not force:
@@ -617,8 +637,8 @@ def smart_rename(
     """
     if smart_isdir(src_path):
         raise IsADirectoryError("%r is a directory" % src_path)
-    src_protocol, _ = SmartPath._extract_protocol(src_path)
-    dst_protocol, _ = SmartPath._extract_protocol(dst_path)
+    src_protocol = SmartPath._extract_protocol(src_path)
+    dst_protocol = SmartPath._extract_protocol(dst_path)
     if src_protocol == dst_protocol:
         SmartPath(src_path).rename(dst_path, overwrite=overwrite)
         return
@@ -634,8 +654,8 @@ def smart_move(src_path: PathLike, dst_path: PathLike, overwrite: bool = True) -
     :param dst_path: Given destination path
     :param overwrite: whether or not overwrite file when exists
     """
-    src_protocol, _ = SmartPath._extract_protocol(src_path)
-    dst_protocol, _ = SmartPath._extract_protocol(dst_path)
+    src_protocol = SmartPath._extract_protocol(src_path)
+    dst_protocol = SmartPath._extract_protocol(dst_path)
     if src_protocol == dst_protocol:
         SmartPath(src_path).rename(dst_path, overwrite=overwrite)
         return
@@ -826,7 +846,7 @@ def _group_glob(globstr: PathLike) -> List[str]:
     expanded_glob = ungloblize(globstr)
 
     for single_glob in expanded_glob:
-        protocol, _ = SmartPath._extract_protocol(single_glob)
+        protocol = SmartPath._extract_protocol(single_glob)
         glob_dict[protocol].append(single_glob)
 
     group_glob_list = []
@@ -1011,7 +1031,7 @@ def smart_load_content(
 def smart_save_content(path: PathLike, content: bytes) -> None:
     """Save bytes content to specified path
 
-    param path: Path to save content
+    :param path: Path to save content
     """
     with smart_open(path, "wb") as fd:
         fd.write(content)
@@ -1021,7 +1041,7 @@ def smart_load_text(path: PathLike) -> str:
     """
     Read content from path
 
-    param path: Path to be read
+    :param path: Path to be read
     """
     with smart_open(path) as fd:
         return fd.read()  # pytype: disable=bad-return-type
@@ -1030,16 +1050,27 @@ def smart_load_text(path: PathLike) -> str:
 def smart_save_text(path: PathLike, text: str) -> None:
     """Save text to specified path
 
-    param path: Path to save text
+    :param path: Path to save text
     """
     with smart_open(path, "w") as fd:
         fd.write(text)
 
 
 class SmartCacher(FileCacher):
+    """Smart cache files in local filesystem"""
+
     cache_path = None
 
     def __init__(self, path: str, cache_path: Optional[str] = None, mode: str = "r"):
+        """
+        :param path: Path to cache
+        :type path: str
+        :param cache_path: Path to cache file, defaults to None, will use ``/tmp``
+        :type cache_path: Optional[str], optional
+        :param mode: Mode to open cache file, defaults to "r"
+        :type mode: str, optional
+        :raises ValueError: If mode is not one of "r", "w", "a"
+        """
         if mode not in ("r", "w", "a"):
             raise ValueError("unacceptable mode: %r" % mode)
         if cache_path is None:
@@ -1059,12 +1090,22 @@ class SmartCacher(FileCacher):
             os.unlink(self.cache_path)
 
 
-def smart_cache(path, cacher=SmartCacher, **options):
+def smart_cache(path, cacher=SmartCacher, **options) -> FileCacher:
     """Return a path to Posixpath Interface
 
-    param path: Path to cache
-    param s3_cacher: Cacher for s3 path
-    param options: Optional arguments for s3_cacher
+    Examples: ::
+
+        >>> import subprocess
+        >>> from megfile import smart_cache
+        >>> with smart_cache(
+        ...     's3://mybucket/myfile.mp4',
+        ...     mode='r',
+        ... ) as cache_path:
+        ...     subprocess.run(['ffprobe', cache_path])
+
+    :param path: Path to cache
+    :param s3_cacher: Cacher for s3 path
+    :param options: Optional arguments for s3_cacher
     """
     if not is_fs(path):
         return cacher(path, **options)
@@ -1074,7 +1115,7 @@ def smart_cache(path, cacher=SmartCacher, **options):
 def smart_touch(path: PathLike):
     """Create a new file on path
 
-    param path: Path to create file
+    :param path: Path to create file
     """
     with smart_open(path, "w"):
         pass
@@ -1083,9 +1124,9 @@ def smart_touch(path: PathLike):
 def smart_getmd5(path: PathLike, recalculate: bool = False, followlinks: bool = False):
     """Get md5 value of file
 
-    param path: File path
-    param recalculate: calculate md5 in real-time or not return s3 etag when path is s3
-    param followlinks: If is True, calculate md5 for real file
+    :param path: File path
+    :param recalculate: calculate md5 in real-time or not return s3 etag when path is s3
+    :param followlinks: If is True, calculate md5 for real file
     """
     return SmartPath(path).md5(recalculate=recalculate, followlinks=followlinks)
 
@@ -1110,9 +1151,9 @@ def smart_concat(src_paths: List[PathLike], dst_path: PathLike) -> None:
     if not src_paths:
         return
 
-    dst_protocol, _ = SmartPath._extract_protocol(dst_path)
+    dst_protocol = SmartPath._extract_protocol(dst_path)
     for src_path in src_paths:
-        src_protocol, _ = SmartPath._extract_protocol(src_path)
+        src_protocol = SmartPath._extract_protocol(src_path)
         if src_protocol != dst_protocol:
             concat_func = _default_concat_func
             break