feat: aliases for SmartPath (#406)

* feat: aliases for SmartPath * make lint happy * update cli & readme * make lint happy 2 * fix review comments * fix review comments 2 --------- Co-authored-by: liyang <[email protected]> Co-authored-by: penghongyang <[email protected]>
megvii-research · Oct 31, 2024 · 300eb9a · 300eb9a
1 parent e92ff9f
commit 300eb9a
Show file tree

Hide file tree

Showing 5 changed files with 149 additions and 33 deletions.
diff --git a/README.md b/README.md
@@ -142,10 +142,10 @@ You can update config file with `megfile` command easyly:
 ```
 $ megfile config s3 accesskey secretkey
 
-# for aliyun
+# for aliyun oss
 $ megfile config s3 accesskey secretkey \
 --addressing-style virtual \
---endpoint-url http://oss-cn-hangzhou.aliyuncs.com \
+--endpoint-url http://oss-cn-hangzhou.aliyuncs.com
 ```
 
 You can get the configuration from `~/.aws/credentials`, like:
@@ -159,6 +159,25 @@ s3 =
     endpoint_url = http://oss-cn-hangzhou.aliyuncs.com
 ```
 
+### Create aliases
+```
+# for volcengine tos
+$ megfile config s3 accesskey secretkey \
+--addressing-style virtual \
+--endpoint-url https://tos-s3-cn-beijing.ivolces.com \
+--profile tos
+
+# create alias
+$ megfile config tos s3+tos
+```
+
+You can get the configuration from `~/.config/megfile/aliases.conf`, like:
+```
+[tos]
+protocol = s3+tos
+```
+
+
 ## How to Contribute
 * We welcome everyone to contribute code to the `megfile` project, but the contributed code needs to meet the following conditions as much as possible:
 

diff --git a/megfile/cli.py b/megfile/cli.py
@@ -570,7 +570,12 @@ def config():
     pass
 
 
-@config.command(short_help="Return the config file for s3")
+def _safe_makedirs(path: str):
+    if path not in ("", ".", "/"):
+        os.makedirs(path, exist_ok=True)
+
+
+@config.command(short_help="Update the config file for s3")
 @click.option(
     "-p",
     "--path",
@@ -584,7 +589,8 @@ def config():
 @click.argument("aws_access_key_id")
 @click.argument("aws_secret_access_key")
 @click.option("-e", "--endpoint-url", help="endpoint-url")
-@click.option("-s", "--addressing-style", help="addressing-style")
+@click.option("-as", "--addressing-style", help="addressing-style")
+@click.option("-sv", "--signature-version", help="signature-version")
 @click.option("--no-cover", is_flag=True, help="Not cover the same-name config")
 def s3(
     path,
@@ -593,6 +599,7 @@ def s3(
     aws_secret_access_key,
     endpoint_url,
     addressing_style,
+    signature_version,
     no_cover,
 ):
     path = os.path.expanduser(path)
@@ -602,30 +609,27 @@ def s3(
         "aws_access_key_id": aws_access_key_id,
         "aws_secret_access_key": aws_secret_access_key,
     }
-    s3 = {}
-    if endpoint_url:
-        s3.update({"endpoint_url": endpoint_url})
-    if addressing_style:
-        s3.update({"addressing_style": addressing_style})
-    if s3:
-        config_dict.update({"s3": s3})
+    s3_config_dict = {
+        "endpoint_url": endpoint_url,
+        "addressing_style": addressing_style,
+        "signature_version": signature_version,
+    }
+
+    s3_config_dict = {k: v for k, v in s3_config_dict.items() if v}
+    if s3_config_dict:
+        config_dict["s3"] = s3_config_dict
 
     def dumps(config_dict: dict) -> str:
         content = "[{}]\n".format(config_dict["name"])
-        content += "aws_access_key_id = {}\n".format(config_dict["aws_access_key_id"])
-        content += "aws_secret_access_key = {}\n".format(
-            config_dict["aws_secret_access_key"]
-        )
+        for key in ("aws_access_key_id", "aws_secret_access_key"):
+            content += "{} = {}\n".format(key, config_dict[key])
         if "s3" in config_dict.keys():
             content += "\ns3 = \n"
-            s3: dict = config_dict["s3"]
-            if "endpoint_url" in s3.keys():
-                content += "    endpoint_url = {}\n".format(s3["endpoint_url"])
-            if "addressing_style" in s3.keys():
-                content += "    addressing_style = {}\n".format(s3["addressing_style"])
+            for key, value in config_dict["s3"].items():
+                content += "    {} = {}\n".format(key, value)
         return content
 
-    os.makedirs(os.path.dirname(path), exist_ok=True)  # make sure dirpath exist
+    _safe_makedirs(os.path.dirname(path))  # make sure dirpath exist
     if not os.path.exists(path):  # If this file doesn't exist.
         content_str = dumps(config_dict)
         with open(path, "w") as fp:
@@ -663,15 +667,15 @@ def dumps(config_dict: dict) -> str:
     click.echo(f"Your oss config has been saved into {path}")
 
 
-@config.command(short_help="Return the config file for s3")
-@click.argument("url")
+@config.command(short_help="Update the config file for hdfs")
 @click.option(
     "-p",
     "--path",
     default="~/.hdfscli.cfg",
-    help="s3 config file, default is $HOME/.hdfscli.cfg",
+    help="hdfs config file, default is $HOME/.hdfscli.cfg",
 )
-@click.option("-n", "--profile-name", default="default", help="s3 config file")
+@click.argument("url")
+@click.option("-n", "--profile-name", default="default", help="hdfs config file")
 @click.option("-u", "--user", help="user name")
 @click.option("-r", "--root", help="hdfs path's root dir")
 @click.option("-t", "--token", help="token for requesting hdfs server")
@@ -681,7 +685,7 @@ def dumps(config_dict: dict) -> str:
     help=f"request hdfs server timeout, default {DEFAULT_HDFS_TIMEOUT}",
 )
 @click.option("--no-cover", is_flag=True, help="Not cover the same-name config")
-def hdfs(url, path, profile_name, user, root, token, timeout, no_cover):
+def hdfs(path, url, profile_name, user, root, token, timeout, no_cover):
     path = os.path.expanduser(path)
     current_config = {
         "url": url,
@@ -704,11 +708,40 @@ def hdfs(url, path, profile_name, user, root, token, timeout, no_cover):
     for key, value in current_config.items():
         if value:
             config[profile_name][key] = value
+
+    _safe_makedirs(os.path.dirname(path))  # make sure dirpath exist
     with open(path, "w") as fp:
         config.write(fp)
     click.echo(f"Your hdfs config has been saved into {path}")
 
 
+@config.command(short_help="Update the config file for aliases")
+@click.option(
+    "-p",
+    "--path",
+    default="~/.config/megfile/aliases.conf",
+    help="alias config file, default is $HOME/.config/megfile/aliases.conf",
+)
+@click.argument("name")
+@click.argument("protocol")
+@click.option("--no-cover", is_flag=True, help="Not cover the same-name config")
+def alias(path, name, protocol, 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,
+    }
+
+    _safe_makedirs(os.path.dirname(path))  # make sure dirpath exist
+    with open(path, "w") as fp:
+        config.write(fp)
+    click.echo(f"Your alias config has been saved into {path}")
+
+
 if __name__ == "__main__":
     # Usage: python -m megfile.cli
     safe_cli()  # pragma: no cover
diff --git a/megfile/smart_path.py b/megfile/smart_path.py
@@ -1,12 +1,17 @@
+import os
+from configparser import ConfigParser
 from pathlib import PurePath
-from typing import Tuple, Union
+from typing import Dict, Tuple, Union
 
 from megfile.lib.compat import fspath
 from megfile.lib.url import get_url_scheme
+from megfile.utils import classproperty
 
 from .errors import ProtocolExistsError, ProtocolNotFoundError
 from .interfaces import BasePath, BaseURIPath, PathLike
 
+aliases_config = "~/.config/megfile/aliases.conf"
+
 
 def _bind_function(name):
     def smart_method(self, *args, **kwargs):
@@ -25,6 +30,17 @@ def smart_property(self):
     return smart_property
 
 
+def _load_aliases_config(config_path) -> Dict[str, Dict[str, str]]:
+    if not os.path.exists(config_path):
+        return {}
+    parser = ConfigParser()
+    parser.read(config_path)
+    configs = {}
+    for section in parser.sections():
+        configs[section] = dict(parser.items(section))
+    return configs
+
+
 class SmartPath(BasePath):
     _registered_protocols = dict()
 
@@ -38,6 +54,13 @@ def __init__(self, path: Union[PathLike, int], *other_paths: PathLike):
             self.path = str(pathlike)
         self.pathlike = pathlike
 
+    @classproperty
+    def _aliases(cls) -> Dict[str, Dict[str, str]]:
+        config_path = os.path.expanduser(aliases_config)
+        aliases = _load_aliases_config(config_path)
+        setattr(cls, "_aliases", aliases)
+        return aliases
+
     @staticmethod
     def _extract_protocol(path: Union[PathLike, int]) -> Tuple[str, Union[str, int]]:
         if isinstance(path, int):
@@ -61,7 +84,11 @@ def _extract_protocol(path: Union[PathLike, int]) -> Tuple[str, Union[str, int]]
 
     @classmethod
     def _create_pathlike(cls, path: Union[PathLike, int]) -> BaseURIPath:
-        protocol, _ = cls._extract_protocol(path)
+        protocol, path_without_protocol = cls._extract_protocol(path)
+        aliases: Dict[str, Dict[str, str]] = cls._aliases  # pyre-ignore[9]
+        if protocol in aliases:
+            protocol = aliases[protocol]["protocol"]
+            path = protocol + "://" + str(path_without_protocol)
         if protocol.startswith("s3+"):
             protocol = "s3"
         if protocol not in cls._registered_protocols:

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -5,6 +5,7 @@
 from click.testing import CliRunner
 
 from megfile.cli import (
+    alias,
     cat,
     cp,
     hdfs,
@@ -355,8 +356,10 @@ def test_config_s3(tmpdir, runner):
             str(tmpdir / "oss_config"),
             "-e",
             "Endpoint",
-            "-s",
-            "Addressing",
+            "-as",
+            "virtual",
+            "-sv",
+            "s3v4",
             "Aws_access_key_id",
             "Aws_secret_access_key",
         ],
@@ -372,7 +375,7 @@ def test_config_s3(tmpdir, runner):
             "new_test",
             "-e",
             "end-point",
-            "-s",
+            "-as",
             "add",
             "1345",
             "2345",
@@ -472,3 +475,20 @@ def test_config_hdfs(tmpdir, runner):
     )
     config.read(str(tmpdir / "config"))
     assert result.exit_code == 1
+
+
+def test_config_alias(tmpdir, runner):
+    result = runner.invoke(
+        alias,
+        [
+            "-p",
+            str(tmpdir / "config"),
+            "a",
+            "b",
+        ],
+    )
+    assert "Your alias config" in result.output
+
+    config = configparser.ConfigParser()
+    config.read(str(tmpdir / "config"))
+    assert config["a"]["protocol"] == "b"
diff --git a/tests/test_smart_path.py b/tests/test_smart_path.py
@@ -3,7 +3,7 @@
 
 import boto3
 import pytest
-from mock import patch
+from mock import PropertyMock, patch
 from moto import mock_aws
 
 from megfile.errors import (
@@ -16,7 +16,7 @@
 from megfile.interfaces import Access
 from megfile.s3_path import S3Path
 from megfile.sftp_path import SftpPath
-from megfile.smart_path import PurePath, SmartPath
+from megfile.smart_path import PurePath, SmartPath, _load_aliases_config, aliases_config
 from megfile.stdio_path import StdioPath
 
 FS_PROTOCOL_PREFIX = FSPath.protocol + "://"
@@ -76,6 +76,23 @@ def test_register_result():
     assert SmartPath.from_uri(FS_TEST_ABSOLUTE_PATH) == SmartPath(FS_TEST_ABSOLUTE_PATH)
 
 
+def test_aliases(fs):
+    config_path = os.path.expanduser(aliases_config)
+    fs.create_file(
+        config_path,
+        contents="[oss2]\nprotocol = s3+oss2\n[tos]\nprotocol = s3+tos",
+    )
+    aliases = {"oss2": {"protocol": "s3+oss2"}, "tos": {"protocol": "s3+tos"}}
+    assert _load_aliases_config(config_path) == aliases
+
+    with patch.object(SmartPath, "_aliases", new_callable=PropertyMock) as mock_aliases:
+        mock_aliases.return_value = aliases
+        assert (
+            SmartPath.from_uri("oss2://bucket/dir/file").pathlike
+            == SmartPath("s3+oss2://bucket/dir/file").pathlike
+        )
+
+
 @patch.object(SmartPath, "_create_pathlike")
 def test_init(funcA):
     SmartPath(FS_TEST_ABSOLUTE_PATH)