megfile 4.1.2__py3-none-any.whl → 4.1.3__py3-none-any.whl

megfile/cli.py

@@ -6,15 +6,19 @@ import sys
 import time
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
+from itertools import islice
 from queue import Queue
 
 import click
+from click import ParamType
+from click.shell_completion import CompletionItem, ZshComplete
 from tqdm import tqdm
 
 from megfile.config import READER_BLOCK_SIZE, SFTP_HOST_KEY_POLICY, set_log_level
 from megfile.hdfs_path import DEFAULT_HDFS_TIMEOUT
 from megfile.interfaces import FileEntry
 from megfile.lib.glob import get_non_glob_dir, has_magic
+from megfile.s3_path import get_s3_session
 from megfile.sftp import sftp_add_host_key
 from megfile.smart import (
     _smart_sync_single_file,
@@ -47,7 +51,6 @@ from megfile.utils import get_human_size
 from megfile.version import VERSION
 
 options = {}
-set_log_level()
 max_file_object_catch_count = 1024 * 128
 
 
@@ -176,8 +179,37 @@ def _ls(path: str, long: bool, recursive: bool, human_readable: bool):
         click.echo(f"total({total_count}): {get_human_size(total_size)}")
 
 
+class PathType(ParamType):
+    name = "path"
+
+    def shell_complete(self, ctx, param, incomplete):
+        if "://" not in incomplete and not incomplete.startswith("/"):
+            completions = [
+                CompletionItem(f"{protocol}://")
+                for protocol in SmartPath._registered_protocols
+            ]
+            for name in get_s3_session().available_profiles:
+                if name == "default":
+                    continue
+                completions.append(CompletionItem(f"s3+{name}://"))
+            return completions
+        try:
+            return [
+                CompletionItem(f"{entry.path}/" if entry.is_dir() else entry.path)
+                for entry in islice(smart_glob_stat(incomplete + "*"), 128)
+            ]
+        except Exception:
+            return []
+
+
+# Some magic, remove trailing spaces in completion
+ZshComplete.source_template = ZshComplete.source_template.replace(
+    "compadd -U -V", "compadd -S '' -U -V"
+)
+
+
 @cli.command(short_help="List all the objects in the path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option(
     "-l",
     "--long",
@@ -202,7 +234,7 @@ def ls(path: str, long: bool, recursive: bool, human_readable: bool):
 
 
 @cli.command(short_help="List all the objects in the path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option(
     "-r",
     "--recursive",
@@ -215,8 +247,8 @@ def ll(path: str, recursive: bool):
 
 
 @cli.command(short_help="Copy files from source to dest, skipping already copied.")
-@click.argument("src_path")
-@click.argument("dst_path")
+@click.argument("src_path", type=PathType())
+@click.argument("dst_path", type=PathType())
 @click.option(
     "-r",
     "--recursive",
@@ -282,8 +314,8 @@ def cp(
 
 
 @cli.command(short_help="Move files from source to dest.")
-@click.argument("src_path")
-@click.argument("dst_path")
+@click.argument("src_path", type=PathType())
+@click.argument("dst_path", type=PathType())
 @click.option(
     "-r",
     "--recursive",
@@ -351,7 +383,7 @@ def mv(
 
 
 @cli.command(short_help="Remove files from path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option(
     "-r",
     "--recursive",
@@ -367,8 +399,8 @@ def rm(path: str, recursive: bool):
 
 
 @cli.command(short_help="Make source and dest identical, modifying destination only.")
-@click.argument("src_path")
-@click.argument("dst_path")
+@click.argument("src_path", type=PathType())
+@click.argument("dst_path", type=PathType())
 @click.option(
     "-f", "--force", is_flag=True, help="Copy files forcible, ignore same files."
 )
@@ -500,7 +532,7 @@ def sync(
 
 
 @cli.command(short_help="Make the path if it doesn't already exist.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def mkdir(path: str):
     _sftp_prompt_host_key(path)
 
@@ -508,7 +540,7 @@ def mkdir(path: str):
 
 
 @cli.command(short_help="Make the file if it doesn't already exist.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def touch(path: str):
     _sftp_prompt_host_key(path)
 
@@ -516,7 +548,7 @@ def touch(path: str):
 
 
 @cli.command(short_help="Concatenate any files and send them to stdout.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def cat(path: str):
     _sftp_prompt_host_key(path)
 
@@ -527,7 +559,7 @@ def cat(path: str):
 @cli.command(
     short_help="Concatenate any files and send first n lines of them to stdout."
 )
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option(
     "-n", "--lines", type=click.INT, default=10, help="print the first NUM lines"
 )
@@ -554,7 +586,7 @@ def _tail_follow_content(path, offset):
 @cli.command(
     short_help="Concatenate any files and send last n lines of them to stdout."
 )
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option(
     "-n", "--lines", type=click.INT, default=10, help="print the last NUM lines"
 )
@@ -601,7 +633,7 @@ def tail(path: str, lines: int, follow: bool):
 
 
 @cli.command(short_help="Write bytes from stdin to file.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 @click.option("-a", "--append", is_flag=True, help="Append to the given file")
 @click.option("-o", "--stdout", is_flag=True, help="File content to standard output")
 def to(path: str, append: bool, stdout: bool):
@@ -626,7 +658,7 @@ def to(path: str, append: bool, stdout: bool):
 
 
 @cli.command(short_help="Produce an md5sum file for all the objects in the path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def md5sum(path: str):
     _sftp_prompt_host_key(path)
 
@@ -634,7 +666,7 @@ def md5sum(path: str):
 
 
 @cli.command(short_help="Return the total size and number of objects in remote:path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def size(path: str):
     _sftp_prompt_host_key(path)
 
@@ -642,7 +674,7 @@ def size(path: str):
 
 
 @cli.command(short_help="Return the mtime and number of objects in remote:path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def mtime(path: str):
     _sftp_prompt_host_key(path)
 
@@ -650,7 +682,7 @@ def mtime(path: str):
 
 
 @cli.command(short_help="Return the stat and number of objects in remote:path.")
-@click.argument("path")
+@click.argument("path", type=PathType())
 def stat(path: str):
     _sftp_prompt_host_key(path)
 
@@ -825,18 +857,26 @@ def hdfs(path, url, profile_name, user, root, token, timeout, no_cover):
     help="alias config file, default is $HOME/.config/megfile/aliases.conf",
 )
 @click.argument("name")
-@click.argument("protocol")
+@click.argument("protocol_or_path")
 @click.option("--no-cover", is_flag=True, help="Not cover the same-name config")
-def alias(path, name, protocol, no_cover):
+def alias(path, name, protocol_or_path, no_cover):
     path = os.path.expanduser(path)
     config = configparser.ConfigParser()
     if os.path.exists(path):
         config.read(path)
     if name in config.sections() and no_cover:
         raise NameError(f"alias-name has been used: {name}")
-    config[name] = {
-        "protocol": protocol,
-    }
+
+    if "://" in protocol_or_path:
+        protocol, prefix = protocol_or_path.split("://", maxsplit=1)
+        config[name] = {
+            "protocol": protocol,
+            "prefix": prefix,
+        }
+    else:
+        config[name] = {
+            "protocol": protocol_or_path,
+        }
 
     _safe_makedirs(os.path.dirname(path))  # make sure dirpath exist
     with open(path, "w") as fp:
@@ -844,6 +884,49 @@ def alias(path, name, protocol, no_cover):
     click.echo(f"Your alias config has been saved into {path}")
 
 
+@cli.group(short_help="Return the completion file")
+def completion():
+    pass
+
+
+@completion.command(short_help="Update the config file for bash")
+def bash():
+    script_name = os.path.basename(sys.argv[0])
+    command = f'eval "$(_{script_name.upper()}_COMPLETE=bash_source {script_name})"'
+    config_path = os.path.expanduser("~/.bashrc")
+    with open(config_path, "r") as fp:
+        if command in fp.read():
+            click.echo("Your bashrc has already been updated.")
+            return
+    with open(config_path, "a") as fp:
+        fp.write("\n" + command + "\n")
+    click.echo("Your bashrc has been updated.")
+
+
+@completion.command(short_help="Update the config file for zsh")
+def zsh():
+    script_name = os.path.basename(sys.argv[0])
+    command = f'eval "$(_{script_name.upper()}_COMPLETE=zsh_source {script_name})"'
+    config_path = os.path.expanduser("~/.zshrc")
+    with open(config_path, "r") as fp:
+        if command in fp.read():
+            click.echo("Your zshrc has already been updated.")
+            return
+    with open(config_path, "a") as fp:
+        fp.write("\n" + command + "\n")
+    click.echo("Your zshrc has been updated.")
+
+
+@completion.command(short_help="Update the config file for fish")
+def fish():
+    script_name = os.path.basename(sys.argv[0])
+    command = f"_{script_name.upper()}_COMPLETE=fish_source {script_name} | source"
+    config_path = os.path.expanduser(f"~/.config/fish/completions/{script_name}.fish")
+    with open(config_path, "w") as fp:
+        fp.write(command)
+    click.echo(f"Your fish config has been saved into {config_path}.")
+
+
 if __name__ == "__main__":
     # Usage: python -m megfile.cli
     safe_cli()  # pragma: no cover

megfile/errors.py

@@ -97,6 +97,7 @@ s3_retry_exceptions = [
     botocore.exceptions.ConnectTimeoutError,
     botocore.exceptions.ProxyConnectionError,
     botocore.exceptions.ConnectionClosedError,
+    botocore.exceptions.SSLError,
     requests.exceptions.ReadTimeout,
     requests.exceptions.ConnectTimeout,
     urllib3.exceptions.IncompleteRead,
@@ -141,8 +142,6 @@ s3_retry_error_codes = (
 def s3_should_retry(error: Exception) -> bool:
     if isinstance(error, s3_retry_exceptions):  # pyre-ignore[6]
         return True
-    if isinstance(error, botocore.exceptions.SSLError):
-        return "EOF" in str(error)
     if isinstance(error, botocore.exceptions.ClientError):
         return client_error_code(error) in s3_retry_error_codes
     return False

megfile/smart_path.py

@@ -81,8 +81,9 @@ class SmartPath(BasePath):
             raise ProtocolNotFoundError("protocol not found: %r" % path)
         aliases: Dict[str, Dict[str, str]] = cls._aliases  # pyre-ignore[9]
         if protocol in aliases:
+            prefix = aliases[protocol].get("prefix", "")
             protocol = aliases[protocol]["protocol"]
-            path = "%s://%s" % (protocol, path_without_protocol)
+            path = "%s://%s%s" % (protocol, prefix, path_without_protocol)
         return protocol, path
 
     @classmethod

megfile/version.py

@@ -1 +1 @@
-VERSION = "4.1.2"
+VERSION = "4.1.3"

{megfile-4.1.2.dist-info → megfile-4.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megfile
-Version: 4.1.2
+Version: 4.1.3
 Summary: Megvii file operation library
 Author-email: megvii <megfile@megvii.com>
 Project-URL: Homepage, https://github.com/megvii-research/megfile

{megfile-4.1.2.dist-info → megfile-4.1.3.dist-info}/RECORD

@@ -1,7 +1,7 @@
 megfile/__init__.py,sha256=7oEfu410CFKzDWZ9RjL5xEJ1gtkJkTfvPrL_7TWdJuY,7366
-megfile/cli.py,sha256=X9nWGke-VM-FCFdRoEoxVH_SNFkKQesBxhSrRv7jAsI,25884
+megfile/cli.py,sha256=iwwlwVNu_yXgqnaURPHqGPrx1WeReqN0zi9C4rqK5Ag,29178
 megfile/config.py,sha256=2MMj5QkhlDJQFZRbCQL2c9iDdeMAVctiaPszRBkg5vM,3988
-megfile/errors.py,sha256=-HuHCMj6AUPlXZnTWlUScZPoiw_DS6xx0NKsSuhztxw,14646
+megfile/errors.py,sha256=A4qX2h1rk6UOOlNy4E81RmI_zuQd5vjh_rlqTfccwag,14589
 megfile/fs.py,sha256=KMEqAE35alpcxiy6du5nPFYcaorhUM_kPJMah3q76ng,19160
 megfile/fs_path.py,sha256=Hozl9LAJ8EMuSWBSZXGj2GNmPZ1sJp9PZs-7hPrLgm8,39341
 megfile/hdfs.py,sha256=owXr4d3j1frCvlbhmhENcSBnKKDky5cJZzWLOF4ZJMo,13251
@@ -15,10 +15,10 @@ megfile/s3_path.py,sha256=zelXhlRVOVSWBE6HJz0vXrrcRzSuj6Cnjd9HLGwPbCM,93644
 megfile/sftp.py,sha256=uBcLQs-j6Q-q-sWAdd-pgi5Qmb_kq7boJM-0sCfcNO0,26540
 megfile/sftp_path.py,sha256=Wz4VcQ0pBUuWDGMSxPpPbutrT09mnY6jZNiAqTi5tO4,43840
 megfile/smart.py,sha256=Sae2KJzaU0k_qV_Bk0YifOMq8WsV5qQ2pGInDRF546I,36411
-megfile/smart_path.py,sha256=HqCOlDwekqqIyJAll-U9YKmaXjjfCGZD5n5aG80lOKw,7592
+megfile/smart_path.py,sha256=Up_6xNZ2019iSzMn_JAU_1H--z-AP6O7SxdXGdeTG0c,7659
 megfile/stdio.py,sha256=ZwxsnJNJYIT7Iyg5pIw4qiyH8bszG6oAhEJuR-hXGG4,658
 megfile/stdio_path.py,sha256=cxaDr8rtisTPnN-rjtaEpqQnshwiqwXFUJBM9xWY7Cg,2711
-megfile/version.py,sha256=kYS0NJJrKL79xx4qm3iMYFp2WaSU6ttQqSdCIGu0WVA,19
+megfile/version.py,sha256=VxmH81mTbRqT3UnqXh5OXAB3IfXRrcnMl0YgLgOqhPw,19
 megfile/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 megfile/lib/base_prefetch_reader.py,sha256=uxVwYknOjc8hLF7q_T2QKMsBqFcrf411ZsuK25CN1eQ,12848
 megfile/lib/combine_reader.py,sha256=Kp2wEloOUpTlIU7dve87MBpSzmIM-F9OtpTawAjFkiU,4828
@@ -43,8 +43,8 @@ megfile/lib/stdio_handler.py,sha256=IDdgENLQlhigEwkLL4zStueVSzdWg7xVcTF_koof_Ek,
 megfile/lib/url.py,sha256=ER32pWy9Q2MAk3TraAaNEBWIqUeBmLuM57ol2cs7-Ks,103
 megfile/utils/__init__.py,sha256=xAzmICA0MtAbg-I2yPfeHjA1N4CiMP4sBrC9BgrfZLw,10151
 megfile/utils/mutex.py,sha256=asb8opGLgK22RiuBJUnfsvB8LnMmodP8KzCVHKmQBWA,2561
-megfile-4.1.2.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-megfile-4.1.2.dist-info/licenses/LICENSE.pyre,sha256=9lf5nT-5ZH25JijpYAequ0bl8E8z5JmZB1qrjiUMp84,1080
+megfile-4.1.3.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+megfile-4.1.3.dist-info/licenses/LICENSE.pyre,sha256=9lf5nT-5ZH25JijpYAequ0bl8E8z5JmZB1qrjiUMp84,1080
 scripts/benchmark/code/iopath_read.py,sha256=O1Qs3mpvD9S_XCuRH2A2FpGWxCOSw6qZvEBrtPeRL1E,661
 scripts/benchmark/code/iopath_write.py,sha256=Mm0efW1J09RJ_CK5i1xtG2hJuaaslikin8qVpuRFP_Q,704
 scripts/benchmark/code/megfile_read.py,sha256=sAMebUiMColHDv3JEkXplImAHvn_IF1-g3BIJxhcQYE,239
@@ -55,8 +55,8 @@ scripts/benchmark/code/s3fs_read.py,sha256=XiTA-qrYblUs-jQWXSnvNg5Wo722C_g47aMMf
 scripts/benchmark/code/s3fs_write.py,sha256=gdXKkWXYGjLJlRT_J64pJN85XvRg3bZexcAJQEMXwtw,402
 scripts/benchmark/code/smart_open_read.py,sha256=SA02jHwS9Y31yFtV9CoJcfND5dR0eA_HsGmGNUrpQls,515
 scripts/benchmark/code/smart_open_write.py,sha256=jDxFJdY97yNH889jz3pawBoei3yaqy8pEMvC_ymHFtM,537
-megfile-4.1.2.dist-info/METADATA,sha256=7Bxhk9yqd8P_kSNqVhvqCP1yif4ScV2x63YE6vzsznQ,9595
-megfile-4.1.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-megfile-4.1.2.dist-info/entry_points.txt,sha256=M6ZWSSv5_5_QtIpZafy3vq7WuOJ_5dSGQQnEZbByt2Q,49
-megfile-4.1.2.dist-info/top_level.txt,sha256=TR6xUw8Px5Ms_UENhEmLNmxOdfHAhTlSt9yTw9LRgsQ,35
-megfile-4.1.2.dist-info/RECORD,,
+megfile-4.1.3.dist-info/METADATA,sha256=hKZKmMsgZAysVrZJ3aRQtiGRrvHnRfEPx68ayMEt3EM,9595
+megfile-4.1.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+megfile-4.1.3.dist-info/entry_points.txt,sha256=M6ZWSSv5_5_QtIpZafy3vq7WuOJ_5dSGQQnEZbByt2Q,49
+megfile-4.1.3.dist-info/top_level.txt,sha256=TR6xUw8Px5Ms_UENhEmLNmxOdfHAhTlSt9yTw9LRgsQ,35
+megfile-4.1.3.dist-info/RECORD,,