YamlLoader.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. # -*- coding: utf-8 -*-
  2. # Author : Charley
  3. # Python : 3.10.8
  4. # Date : 2025/12/22 10:44
  5. import os, re
  6. import yaml
  7. regex = re.compile(r'^\$\{(?P<ENV>[A-Z_\-]+:)?(?P<VAL>[\w.]+)}$')
  8. class YamlConfig:
  9. def __init__(self, config):
  10. self.config = config
  11. def get(self, key: str):
  12. return YamlConfig(self.config.get(key))
  13. def getValueAsString(self, key: str):
  14. try:
  15. match = regex.match(self.config[key])
  16. group = match.groupdict()
  17. if group['ENV'] is not None:
  18. env = group['ENV'][:-1]
  19. return os.getenv(env, group['VAL'])
  20. return None
  21. except:
  22. return self.config[key]
  23. def getValueAsInt(self, key: str):
  24. try:
  25. match = regex.match(self.config[key])
  26. group = match.groupdict()
  27. if group['ENV'] is not None:
  28. env = group['ENV'][:-1]
  29. return int(os.getenv(env, group['VAL']))
  30. return 0
  31. except:
  32. return int(self.config[key])
  33. def getValueAsBool(self, key: str):
  34. try:
  35. match = regex.match(self.config[key])
  36. group = match.groupdict()
  37. if group['ENV'] is not None:
  38. env = group['ENV'][:-1]
  39. return bool(os.getenv(env, group['VAL']))
  40. return False
  41. except:
  42. return bool(self.config[key])
  43. def _resolve_path(path: str) -> str:
  44. """
  45. 解析 yaml 文件路径,按优先级查找:
  46. 1) 绝对路径或 cwd 下存在 → 直接用(保留旧行为,向后兼容)
  47. 2) 调用方主脚本所在目录 → 兜底,方便打包后从任意 cwd 启动
  48. :param path: (str) 用户传入的路径,默认 'application.yml'
  49. :return: (str) 实际可读取的完整路径;找不到则返回原 path 让 open() 抛错
  50. """
  51. # 1) 旧行为:cwd 或绝对路径
  52. if os.path.exists(path):
  53. return path
  54. # 2) 主脚本目录(__main__.__file__)
  55. try:
  56. import __main__
  57. main_file = getattr(__main__, '__file__', None)
  58. if main_file:
  59. candidate = os.path.join(os.path.dirname(os.path.abspath(main_file)), path)
  60. if os.path.exists(candidate):
  61. return candidate
  62. except Exception:
  63. pass
  64. return path
  65. def readYaml(path: str = 'application.yml', profile: str = None) -> YamlConfig:
  66. """
  67. 读取 yaml 配置。
  68. :param path: (str) yaml 文件路径,默认 'application.yml'。
  69. 优先 cwd / 绝对路径(保留旧行为),找不到再 fallback 到主脚本所在目录。
  70. :param profile: (str) 可选环境后缀,如 'dev' 会额外加载 'application-dev.yml' 并 update
  71. :return: (YamlConfig) 配置访问对象
  72. :raises FileNotFoundError: cwd 和主脚本目录都找不到时抛出
  73. """
  74. real_path = _resolve_path(path)
  75. with open(real_path, encoding='utf-8') as fd:
  76. conf = yaml.load(fd, Loader=yaml.FullLoader)
  77. if profile is not None:
  78. result = real_path.rsplit('.', 1)
  79. profiledYaml = f'{result[0]}-{profile}.{result[1]}'
  80. if os.path.exists(profiledYaml):
  81. with open(profiledYaml, encoding='utf-8') as fd:
  82. conf.update(yaml.load(fd, Loader=yaml.FullLoader))
  83. return YamlConfig(conf)