feat: host params for login cli

Author: SAKURA-CAT

swanlab/cli/commands/auth/login.py

@@ -7,11 +7,15 @@
 @Description:
     登录模块
 """
+import os
+import sys
+
 import click
 from swankit.log import FONT
 
+from swanlab import SwanLabEnv
 from swanlab.api.auth import terminal_login
-from swanlab.package import has_api_key
+from swanlab.package import has_api_key, HostFormatter
 
 
 @click.command()
@@ -30,13 +34,36 @@
     help="If you prefer not to engage in commands-line interaction to input the api key, "
     "this will allow automatic login.",
 )
-def login(api_key: str, relogin: bool):
+@click.option(
+    "--host",
+    "-h",
+    default=None,
+    type=str,
+    help="The host of the swanlab server.",
+)
+@click.option(
+    "--web-host",
+    "-w",
+    default=None,
+    type=str,
+    help="The web host of the swanlab cloud front-end.",
+)
+def login(api_key: str, relogin: bool, host: str, web_host: str):
     """Login to the swanlab cloud."""
     if not relogin and has_api_key():
         # 此时代表token已经获取，需要打印一条信息：已经登录
         command = FONT.bold("swanlab login --relogin")
         tip = FONT.swanlab("You are already logged in. Use `" + command + "` to force relogin.")
         return print(tip)
+    # 清除环境变量
+    if relogin:
+        del os.environ[SwanLabEnv.API_HOST.value]
+        del os.environ[SwanLabEnv.WEB_HOST.value]
+    try:
+        HostFormatter(host, web_host)()
+    except ValueError as e:
+        click.BadParameter(str(e))
+        return sys.exit(1)
     # 进行登录，此时将直接覆盖本地token文件
     login_info = terminal_login(api_key)
     print(FONT.swanlab("Login successfully. Hi, " + FONT.bold(FONT.default(login_info.username))) + "!")

swanlab/data/sdk.py

@@ -8,7 +8,6 @@
     在此处封装swanlab在日志记录模式下的各种接口
 """
 import os
-import re
 from typing import Optional, Union, Dict, Tuple, Literal
 
 from swanboard import SwanBoardCallback
@@ -30,7 +29,7 @@
 )
 from .run.helper import SwanLabRunOperator
 from ..error import KeyFileError
-from ..package import get_key, get_host_web
+from ..package import get_key, get_host_web, HostFormatter
 
 
 def _check_proj_name(name: str) -> str:
@@ -58,31 +57,6 @@ def _check_proj_name(name: str) -> str:
     return _name
 
 
-class HostFormatter:
-    def __init__(self):
-        # 定义正则模式，匹配协议、主机、端口三部分
-        self.pattern = re.compile(
-            r'^(?:(https?)://)?'  # 可选协议 http 或 https
-            r'([a-zA-Z0-9.-]+\.[a-zA-Z]{2,63})'  # 必填 主机名必须包含顶级域名（TLD）
-            r'(?::(\d{1,5}))?$'  # 可选端口号（1~5位数字）
-        )
-
-    def __call__(self, input_str: str) -> str:
-        match = self.pattern.match(input_str.rstrip("/"))
-        if match:
-            protocol = match.group(1) or "https"  # 默认协议为 https
-            host = match.group(2)
-            port = match.group(3)
-
-            # 构建标准化的 URL 输出
-            result = f"{protocol}://{host}"
-            if port:
-                result += f":{port}"
-            return result
-        else:
-            raise ValueError("Invalid host format")
-
-
 def login(api_key: str = None, host: str = None, web_host: str = None):
     """
     Login to SwanLab Cloud. If you already have logged in, you can use this function to relogin.
@@ -100,19 +74,9 @@ def login(api_key: str = None, host: str = None, web_host: str = None):
     """
     if SwanLabRun.is_started():
         raise RuntimeError("You must call swanlab.login() before using init()")
-    # ---------------------------------- 检查host是否合法，并格式化 ----------------------------------
-    formater = HostFormatter()
-    if host:
-        try:
-            os.environ[SwanLabEnv.API_HOST.value] = formater(host) + "/api"
-        except ValueError:
-            raise ValueError("Invalid host: {}".format(host))
-    if web_host:
-        try:
-            os.environ[SwanLabEnv.WEB_HOST.value] = formater(web_host)
-        except ValueError:
-            raise ValueError("Invalid web_host: {}".format(web_host))
-    # ---------------------------------- 登录，初始化http对象 ----------------------------------
+    # 检查host是否合法，并格式化，注入到环境变量中
+    HostFormatter(host, web_host)()
+    # 登录，初始化http对象
     CloudRunCallback.login_info = code_login(api_key) if api_key else CloudRunCallback.create_login_info()
 
 

swanlab/package.py

@@ -10,6 +10,7 @@
 import json
 import netrc
 import os
+import re
 from typing import Optional
 
 import requests
@@ -54,6 +55,50 @@ def get_package_latest_version(timeout=0.5) -> Optional[str]:
 # ---------------------------------- 云端相关 ----------------------------------
 
 
+class HostFormatter:
+    def __init__(self, host: str = None, web_host: str = None):
+        # 定义正则模式，匹配协议、主机、端口三部分
+        self.pattern = re.compile(
+            r'^(?:(https?)://)?'  # 可选协议 http 或 https
+            r'([a-zA-Z0-9.-]+\.[a-zA-Z]{2,63})'  # 必填 主机名必须包含顶级域名（TLD）
+            r'(?::(\d{1,5}))?$'  # 可选端口号（1~5位数字）
+        )
+        self.host = host
+        self.web_host = web_host
+
+    def fmt(self, input_str: str) -> str:
+        match = self.pattern.match(input_str.rstrip("/"))
+        if match:
+            protocol = match.group(1) or "https"  # 默认协议为 https
+            host = match.group(2)
+            port = match.group(3)
+
+            # 构建标准化的 URL 输出
+            result = f"{protocol}://{host}"
+            if port:
+                result += f":{port}"
+            return result
+        else:
+            raise ValueError("Invalid host format")
+
+    def __call__(self):
+        """
+        如果host或web_host不为空，格式化并设置环境变量
+        :raises ValueError: host或web_host格式不正确
+        """
+        if self.host:
+            try:
+                os.environ[SwanLabEnv.API_HOST.value] = self.fmt(self.host) + "/api"
+            except ValueError:
+                raise ValueError("Invalid host: {}".format(self.host))
+            self.web_host = self.host
+        if self.web_host:
+            try:
+                os.environ[SwanLabEnv.WEB_HOST.value] = self.fmt(self.web_host)
+            except ValueError:
+                raise ValueError("Invalid web_host: {}".format(self.web_host))
+
+
 def get_host_web() -> str:
     """获取swanlab网站网址
     :return: swanlab网站的网址

test/unit/cli/test_cli_login.py

@@ -7,13 +7,17 @@
 @Description:
     测试命令登录
 """
+import os
+
 import nanoid
-from swanlab.package import get_key
+import pytest
 from click.testing import CliRunner
+
+import tutils as T
 from swanlab.cli.main import cli
+from swanlab.env import SwanLabEnv
 from swanlab.error import ValidationError, APIKeyFormatError
-import tutils as T
-import pytest
+from swanlab.package import get_key
 
 
 # noinspection PyTypeChecker
@@ -35,3 +39,20 @@ def test_login_fail():
     result = runner.invoke(cli, ["login", "--api-key", "wrong-key"])
     assert result.exit_code == 1
     assert isinstance(result.exception, APIKeyFormatError)
+
+
+# noinspection PyTypeChecker
+@pytest.mark.skipif(T.is_skip_cloud_test, reason="skip cloud test")
+def test_login_host():
+    """
+    测试登录时指定host
+    """
+    runner = CliRunner()
+    del os.environ[SwanLabEnv.API_HOST.value]  # 删除环境变量
+    del os.environ[SwanLabEnv.WEB_HOST.value]  # 删除环境变量
+    result = runner.invoke(cli, ["login", "--api-key", T.API_KEY, "--host", T.API_HOST.rstrip("/api")])
+    assert result.exit_code == 0
+    del os.environ[SwanLabEnv.API_HOST.value]  # 删除环境变量
+    del os.environ[SwanLabEnv.WEB_HOST.value]  # 删除环境变量
+    result = runner.invoke(cli, ["login", "--api-key", T.API_KEY, "--host", "http://wrong-host"])
+    assert result.exit_code == 1

test/unit/data/test_sdk.py

@@ -306,25 +306,6 @@ def test_init_logdir_env(self):
         assert run.public.swanlog_dir == logdir
 
 
-class TestHostFormatter:
-    def test_ok(self):
-        formatter = S.HostFormatter()
-        assert formatter("swanlab.cn") == "https://swanlab.cn"
-        assert formatter("https://swanlab.cn") == "https://swanlab.cn"
-        assert formatter("http://swanlab.cn") == "http://swanlab.cn"  # noqa
-        assert formatter("https://swanlab.cn:8443/") == "https://swanlab.cn:8443"
-        assert formatter("abc.example.com") == "https://abc.example.com"
-
-    def test_value_err(self):
-        formatter = S.HostFormatter()
-        with pytest.raises(ValueError):
-            formatter("test")
-        with pytest.raises(ValueError):
-            formatter("https://test")
-        with pytest.raises(ValueError):
-            formatter("http://test")  # noqa
-
-
 @pytest.mark.skipif(T.is_skip_cloud_test, reason="skip cloud test")
 class TestLogin:
     """

test/unit/test_package.py

@@ -239,3 +239,22 @@ def test_wrong_host(self):
         self.save_api_key()
         os.environ[SwanLabEnv.API_HOST.value] = nanoid.generate()
         assert not P.has_api_key()
+
+
+class TestHostFormatter:
+    def test_ok(self):
+        formatter = P.HostFormatter()
+        assert formatter.fmt("swanlab.cn") == "https://swanlab.cn"
+        assert formatter.fmt("https://swanlab.cn") == "https://swanlab.cn"
+        assert formatter.fmt("http://swanlab.cn") == "http://swanlab.cn"  # noqa
+        assert formatter.fmt("https://swanlab.cn:8443/") == "https://swanlab.cn:8443"
+        assert formatter.fmt("abc.example.com") == "https://abc.example.com"
+
+    def test_value_err(self):
+        formatter = P.HostFormatter()
+        with pytest.raises(ValueError):
+            formatter.fmt("test")
+        with pytest.raises(ValueError):
+            formatter.fmt("https://test")
+        with pytest.raises(ValueError):
+            formatter.fmt("http://test")  # noqa