From 518623186933193d1279a6bf9ccb8c1aa2ee600a Mon Sep 17 00:00:00 2001
From: yanyongyu <yanyongyu_1@126.com>
Date: Fri, 7 Aug 2020 11:18:02 +0800
Subject: [PATCH] change config load method to support custom configs

---
 nonebot/__init__.py |  2 +-
 nonebot/config.py   | 81 +++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/nonebot/__init__.py b/nonebot/__init__.py
index ab5ac5a8..cef2f0a0 100644
--- a/nonebot/__init__.py
+++ b/nonebot/__init__.py
@@ -29,7 +29,7 @@ def get_asgi():
     return driver.asgi
 
 
-def init(*, _env_file=None, **kwargs):
+def init(*, _env_file: Optional[str] = None, **kwargs):
     global _driver
     env = Env()
     config = Config(**kwargs, _env_file=_env_file or f".env.{env.environment}")
diff --git a/nonebot/config.py b/nonebot/config.py
index a58c0dea..1e8da0ec 100644
--- a/nonebot/config.py
+++ b/nonebot/config.py
@@ -1,11 +1,72 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+import os
+from pathlib import Path
 from datetime import timedelta
 from ipaddress import IPv4Address
-from typing import Set, Dict, Union, Optional
+from typing import Set, Dict, Union, Mapping, Optional
 
 from pydantic import BaseSettings
+from pydantic.env_settings import SettingsError, env_file_sentinel, read_env_file
+
+
+class BaseConfig(BaseSettings):
+
+    def _build_environ(
+            self,
+            _env_file: Union[Path, str,
+                             None] = None) -> Dict[str, Optional[str]]:
+        """
+        Build environment variables suitable for passing to the Model.
+        """
+        d: Dict[str, Optional[str]] = {}
+
+        if self.__config__.case_sensitive:
+            env_vars: Mapping[str, Optional[str]] = os.environ
+        else:
+            env_vars = {k.lower(): v for k, v in os.environ.items()}
+
+        env_file = _env_file if _env_file != env_file_sentinel else self.__config__.env_file
+        env_file_vars = {}
+        if env_file is not None:
+            env_path = Path(env_file)
+            if env_path.is_file():
+                env_file_vars = read_env_file(
+                    env_path, case_sensitive=self.__config__.case_sensitive)
+                env_vars = {**env_file_vars, **env_vars}
+
+        for field in self.__fields__.values():
+            env_val: Optional[str] = None
+            for env_name in field.field_info.extra['env_names']:
+                env_val = env_vars.get(env_name)
+                if env_name in env_file_vars:
+                    del env_file_vars[env_name]
+                if env_val is not None:
+                    break
+
+            if env_val is None:
+                continue
+
+            if field.is_complex():
+                try:
+                    env_val = self.__config__.json_loads(env_val)
+                except ValueError as e:
+                    raise SettingsError(
+                        f'error parsing JSON for "{env_name}"'  # type: ignore
+                    ) from e
+            d[field.alias] = env_val
+
+        if env_file_vars:
+            for env_name, env_val in env_file_vars.items():
+                try:
+                    env_val = self.__config__.json_loads(env_val)
+                except ValueError as e:
+                    pass
+
+                d[env_name] = env_val
+
+        return d
 
 
 class Env(BaseSettings):
@@ -15,7 +76,19 @@ class Env(BaseSettings):
         env_file = ".env"
 
 
-class Config(BaseSettings):
+class Config(BaseConfig):
+    """
+    NoneBot Config Object
+
+    configs:
+
+    ### `driver`
+
+    - 类型: `str`
+    - 默认值: `"nonebot.drivers.fastapi"`
+    - 说明:
+      nonebot 运行使用后端框架封装 Driver 。继承自 nonebot.driver.BaseDriver 。
+    """
     # nonebot configs
     driver: str = "nonebot.drivers.fastapi"
     host: IPv4Address = IPv4Address("127.0.0.1")
@@ -33,7 +106,9 @@ class Config(BaseSettings):
     session_expire_timeout: timedelta = timedelta(minutes=2)
 
     # custom configs
-    custom_config: dict = {}
+    # custom configs can be assigned during nonebot.init
+    # or from env file using json loads
 
     class Config:
+        extra = "allow"
         env_file = ".env.prod"