fix: support disable rich console (#38)

feat: support disable rich console

Signed-off-by: Keming <kemingyang@tensorchord.ai>

modelz/client.py

@@ -5,13 +5,12 @@
 from urllib.parse import urljoin
 
 import httpx
-from rich.console import Console
 
+from modelz.console import console
 from modelz.env import EnvConfig
 from modelz.serde import Serde, SerdeEnum, TextSerde
 
 TIMEOUT = httpx.Timeout(5, read=300, write=300)
-console = Console()
 config = EnvConfig()
 DEFAULT_RESP_SERDE = TextSerde()
 DEFAULT_RETRY = 3

modelz/cmd.py

@@ -1,12 +1,9 @@
 import sys
 from typing import Dict, Optional
 
-from rich.console import Console
-
 from modelz.args import build_argument_parser, parse_arguments
 from modelz.client import ModelzClient
-
-console = Console()
+from modelz.console import console
 
 
 def inference(
@@ -29,7 +26,6 @@ def inference(
         write_file: write response to the file
         key: API key, will try to read from env ``MODELZ_API_KEY`` if not provided
         serde: serilize/deserilize method, choose from [json|msgpack|raw]
-        output: output target, choose from [console|file]
     """
     data = params
     if read_stdin:

modelz/console.py

@@ -0,0 +1,32 @@
+import re
+
+from rich.console import Console
+
+from modelz.env import EnvConfig
+
+config = EnvConfig()
+
+
+class ModelzConsole:
+    class Status:
+        def __init__(self, msg) -> None:
+            self.msg = msg
+
+        def __enter__(self):
+            print(self.msg)
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+    def __init__(self) -> None:
+        self._remove_style = re.compile(r"\[.*?\]")
+        self.formatter = lambda msg: self._remove_style.sub("", msg)
+
+    def status(self, msg: str):
+        return self.Status(self.formatter(msg))
+
+    def print(self, msg: str):
+        print(self.formatter(msg))
+
+
+console = ModelzConsole() if config.disable_rich else Console()

modelz/env.py

@@ -9,10 +9,12 @@ class EnvConfig:
     def __init__(self) -> None:
         self.api_key: str
         self.host: str = "https://{}.modelz.io/"
+        self.disable_rich: bool = True
         self.update_from_env()
 
     def update_from_env(self):
-        for key in ("api_key", "host"):
+        for key in ("api_key", "host", "disable_rich"):
             val = os.environ.get(f"{PREFIX}{key.upper()}")
             if val is not None:
-                setattr(self, key, val)
+                setattr(self, key, type(key)(val))
+                print(f"update {key} from env: {val}")