feat: alias support prefix

megvii-research · Dec 7, 2024 · 4fe2cb3 · 4fe2cb3
1 parent 81d9691
commit 4fe2cb3
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 8 deletions.
diff --git a/megfile/cli.py b/megfile/cli.py
@@ -734,18 +734,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:

diff --git a/megfile/smart_path.py b/megfile/smart_path.py
@@ -81,8 +81,9 @@ def _extract_protocol(
             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

diff --git a/tests/test_smart_path.py b/tests/test_smart_path.py
@@ -19,6 +19,8 @@
 from megfile.smart_path import PurePath, SmartPath, _load_aliases_config, aliases_config
 from megfile.stdio_path import StdioPath
 
+from .test_sftp import sftp_mocker  # noqa: F401
+
 FS_PROTOCOL_PREFIX = FSPath.protocol + "://"
 FS_TEST_ABSOLUTE_PATH = "/test/dir/file"
 FS_TEST_ABSOLUTE_PATH_WITH_PROTOCOL = FS_PROTOCOL_PREFIX + FS_TEST_ABSOLUTE_PATH
@@ -76,7 +78,7 @@ def test_register_result():
     assert SmartPath.from_uri(FS_TEST_ABSOLUTE_PATH) == SmartPath(FS_TEST_ABSOLUTE_PATH)
 
 
-def test_aliases(fs):
+def test_aliases(fs, sftp_mocker):
     config_path = os.path.expanduser(aliases_config)
     fs.create_file(
         config_path,
@@ -88,10 +90,18 @@ def test_aliases(fs):
     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("oss2://bucket/dir/file").pathlike
             == SmartPath("s3+oss2://bucket/dir/file").pathlike
         )
 
+    aliases = {"dev": {"protocol": "sftp", "prefix": "ubuntu@host//"}}
+    with patch.object(SmartPath, "_aliases", new_callable=PropertyMock) as mock_aliases:
+        mock_aliases.return_value = aliases
+        assert (
+            SmartPath("dev://dir/file").pathlike
+            == SmartPath("sftp://ubuntu@host//dir/file").pathlike
+        )
+
 
 @patch.object(SmartPath, "_create_pathlike")
 def test_init(funcA):