zhang_crawl_bio/ui/backend/utils/config.py

import yaml
from pathlib import Path
from pydantic import BaseModel, Field
from typing import List, Dict, Union,Optional,Any
from utils.pydantic_auto_field import AutoLoadModel
APP_PATH = Path(__file__).parent.parent
OUTPUT_DIR = APP_PATH / "output"
CONFIG_PATH = APP_PATH / "config.yaml"
REPO_BASE_DIR = Path(APP_PATH.parent.parent)
DOWNLOAD_DIR = REPO_BASE_DIR / "download"
PROXY_POLL_DIR = DOWNLOAD_DIR / "proxy_pool"
WORKER_DIR_BASE = APP_PATH.parent.parent

class Proxy(BaseModel):
    name: Optional[str] = ''
    port: Optional[int] = -1
    file_path: Optional[str] = ''
    startup: Optional[bool] = False

class Sub(AutoLoadModel):
    url: Optional[str] = None
    start_port: Optional[int] = 9660
    redis_url: Optional[str] = 'redis://localhost:6379/8'
    file: Optional[str] = str(REPO_BASE_DIR / r'download\proxy_pool\6137e542.yaml')
    temp_dir: Optional[str] = str(PROXY_POLL_DIR / "temp")  
    auto_start: Optional[bool] = True 
    proxies: Optional[Dict[Union[int,str], Proxy]] = {}
    
    def __init__(self, **data):
        super().__init__(**data)
        if self.proxies:
            self.proxies = {k: Proxy(**v) if isinstance(v, dict) else v for k, v in self.proxies.items()}

class Browser(BaseModel):
    exe_path: Optional[str] = str(REPO_BASE_DIR / r"download\GoogleChromePortable\GoogleChromePortable.exe")
    no_imgs: Optional[bool] = True

class Backend(BaseModel):
    host: Optional[str] = "localhost"
    port: Optional[int] = 5835

class RedisConfig(BaseModel):
    host: Optional[str] = "localhost"
    port: Optional[int] = 6379
    db: Optional[int] = 1
    exe: Optional[str] = str(REPO_BASE_DIR / r"download\Redis-x64-5.0.14.1\redis-server.exe")

class Config(BaseModel):
    sub: Optional[Sub] = Sub()
    select_proxy: Optional[str] = "system"
    mimo_exe: Optional[str] = str(PROXY_POLL_DIR / r"mihomo-windows-amd64-go120.exe")
    pandoc_exe: Optional[str] = str(REPO_BASE_DIR / r"download\pandoc-3.6.3-windows-x86_64\pandoc.exe")
    worker_backend_py: Optional[str] = str(WORKER_DIR_BASE / r"worker\api\worker_server.py")
    sqluri: Optional[str] = 'sqlite:///' + str(REPO_BASE_DIR / r'output\temp.db')
    browser: Optional[Browser] = Browser()
    backend: Optional[Backend] = Backend()
    worker_backend: Optional[Backend] = Field(default_factory=lambda: Backend(host="localhost", port=8003))
    redis: Optional[RedisConfig] = RedisConfig()

    def save(self):
        config_path = get_config_path()
        with open(config_path, "w", encoding="utf-8") as file:
            yaml.dump(self.model_dump(), file)
        return self
            
def get_config_path():
    return CONFIG_PATH

def read_config(config_path: Path):
    if not config_path.exists():
        config = Config()
        config.save()
        return config
    with open(config_path, "r", encoding="utf-8") as file:
        config_dict = yaml.safe_load(file)
    return Config(**config_dict)

config = read_config(get_config_path())
WORKER_SERVICE_URL = f"http://{config.worker_backend.host}:{config.worker_backend.port}"

def main():
    print(config)
    config.save()

if __name__ == "__main__":
    main()