Refactor remote api config to config.py

ai_models/remote/api.py

@@ -8,6 +8,8 @@
 from multiurl import download, robust
 from tqdm import tqdm
 
+from .config import API_URL, CONFIG_PATH, load_config
+
 LOG = logging.getLogger(__name__)
 
 
@@ -28,37 +30,23 @@ def __init__(
         url: str = None,
         token: str = None,
     ):
-        root = os.path.join(os.path.expanduser("~"), ".config", "ai-models")
-        os.makedirs(root, exist_ok=True)
-
-        configfile = os.path.join(root, "api.yaml")
-
-        if os.path.exists(configfile):
-            from yaml import safe_load
-
-            with open(configfile, "r") as f:
-                config = safe_load(f) or {}
+        config = load_config()
 
-        url = (
-            url
-            or os.getenv("AI_MODELS_REMOTE_URL")
-            or config.get("url")
-            or "https://ai-models.ecmwf.int"
+        self.url = (
+            url or os.getenv("AI_MODELS_REMOTE_URL") or config.get("url") or API_URL
         )
-        LOG.info("Using remote server %s", url)
+        self.token = token or os.getenv("AI_MODELS_REMOTE_TOKEN") or config.get("token")
 
-        token = token or os.getenv("AI_MODELS_REMOTE_TOKEN") or config.get("token")
-
-        if not token:
+        if not self.token:
             LOG.error(
-                "Missing remote token. Set it in %s or in env AI_MODELS_REMOTE_TOKEN",
-                configfile,
+                "Missing remote token. Set it in %s or env AI_MODELS_REMOTE_TOKEN",
+                CONFIG_PATH,
             )
             sys.exit(1)
 
-        self.url = url
-        self.token = token
-        self.auth = BearerAuth(token)
+        LOG.info("Using remote server %s", self.url)
+
+        self.auth = BearerAuth(self.token)
         self.output_file = output_file
         self.input_file = input_file
         self._timeout = 300

ai_models/remote/config.py

@@ -0,0 +1,42 @@
+import logging
+import os
+
+API_URL = "https://ai-models.ecmwf.int"
+
+ROOT_PATH = os.path.join(os.path.expanduser("~"), ".config", "ai-models")
+CONFIG_PATH = os.path.join(ROOT_PATH, "api.yaml")
+
+LOG = logging.getLogger(__name__)
+
+
+def config_exists():
+    return os.path.exists(CONFIG_PATH)
+
+
+def create_config():
+    if config_exists():
+        return
+
+    try:
+        os.makedirs(ROOT_PATH, exist_ok=True)
+        with open(CONFIG_PATH, "w") as f:
+            f.write("token: \n")
+            f.write(f"url: {API_URL}\n")
+    except Exception as e:
+        LOG.error(f"Failed to create config {CONFIG_PATH}")
+        LOG.error(e, exc_info=True)
+
+
+def load_config() -> dict:
+    from yaml import safe_load
+
+    if not config_exists():
+        create_config()
+
+    try:
+        with open(CONFIG_PATH, "r") as f:
+            return safe_load(f) or {}
+    except Exception as e:
+        LOG.error(f"Failed to read config {CONFIG_PATH}")
+        LOG.error(e, exc_info=True)
+        return {}