config.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189
  1. #!/usr/bin/env python3
  2. """
  3. 配置模块 - 统一管理项目配置
  4. """
  5. import os
  6. from pathlib import Path
  7. from typing import Optional
  8. class Config:
  9. """项目配置类"""
  10. # 默认缓存根目录(用户主目录下的 cache)
  11. _DEFAULT_CACHE_ROOT = os.path.expanduser("~/cache")
  12. # 缓存根目录
  13. _cache_root: Optional[str] = None
  14. @classmethod
  15. def get_cache_root(cls) -> str:
  16. """
  17. 获取缓存根目录
  18. Returns:
  19. 缓存根目录路径
  20. """
  21. if cls._cache_root is None:
  22. # 1. 优先从环境变量读取
  23. cache_root = os.environ.get("CACHE_ROOT")
  24. if cache_root:
  25. cls._cache_root = cache_root
  26. else:
  27. # 2. 使用默认路径
  28. cls._cache_root = cls._DEFAULT_CACHE_ROOT
  29. return cls._cache_root
  30. @classmethod
  31. def set_cache_root(cls, path: str) -> None:
  32. """
  33. 设置缓存根目录
  34. Args:
  35. path: 缓存根目录路径(可以是绝对路径或相对路径)
  36. """
  37. cls._cache_root = path
  38. @classmethod
  39. def get_cache_dir(cls, subdir: str) -> str:
  40. """
  41. 获取特定子模块的缓存目录
  42. Args:
  43. subdir: 子目录名称,如:
  44. - "text_embedding", "semantic_similarity" - 计算缓存
  45. - "data/search", "data/detail" - 爬虫数据缓存
  46. - "data/analysis" - 分析结果缓存
  47. Returns:
  48. 完整的缓存目录路径
  49. """
  50. cache_root = cls.get_cache_root()
  51. return str(Path(cache_root) / subdir)
  52. @classmethod
  53. def get_data_dir(cls, subdir: str = "") -> str:
  54. """
  55. 获取数据缓存目录(data 目录现在在缓存根目录下)
  56. Args:
  57. subdir: 子目录名称,如 "search", "detail", "tools_list" 等
  58. 如果为空字符串,返回 data 根目录
  59. Returns:
  60. 完整的数据目录路径
  61. Note:
  62. data 目录现在统一放在缓存根目录下:
  63. - 默认:cache/data/
  64. - 如果设置了 CACHE_ROOT=/custom: /custom/data/
  65. """
  66. cache_root = cls.get_cache_root()
  67. if subdir:
  68. return str(Path(cache_root) / "data" / subdir)
  69. return str(Path(cache_root) / "data")
  70. @classmethod
  71. def reset(cls) -> None:
  72. """
  73. 重置配置为默认值(主要用于测试)
  74. """
  75. cls._cache_root = None
  76. # 便捷函数
  77. def get_cache_root() -> str:
  78. """获取缓存根目录"""
  79. return Config.get_cache_root()
  80. def set_cache_root(path: str) -> None:
  81. """设置缓存根目录"""
  82. Config.set_cache_root(path)
  83. def get_cache_dir(subdir: str) -> str:
  84. """获取特定子模块的缓存目录"""
  85. return Config.get_cache_dir(subdir)
  86. def get_data_dir(subdir: str = "") -> str:
  87. """
  88. 获取数据缓存目录
  89. Note: data 目录现在在缓存根目录下,例如 cache/data/
  90. """
  91. return Config.get_data_dir(subdir)
  92. if __name__ == "__main__":
  93. print("=" * 60)
  94. print("配置模块示例")
  95. print("=" * 60)
  96. print()
  97. # 示例 1: 使用默认配置
  98. print("示例 1: 默认配置")
  99. print(f"缓存根目录: {get_cache_root()}")
  100. print(f"text_embedding 缓存: {get_cache_dir('text_embedding')}")
  101. print(f"semantic_similarity 缓存: {get_cache_dir('semantic_similarity')}")
  102. print()
  103. # 示例 2: 自定义缓存根目录
  104. print("示例 2: 自定义缓存根目录")
  105. set_cache_root("/tmp/my_cache")
  106. print(f"缓存根目录: {get_cache_root()}")
  107. print(f"text_embedding 缓存: {get_cache_dir('text_embedding')}")
  108. print(f"semantic_similarity 缓存: {get_cache_dir('semantic_similarity')}")
  109. print()
  110. # 示例 3: 使用相对路径
  111. print("示例 3: 使用相对路径")
  112. set_cache_root("data/cache")
  113. print(f"缓存根目录: {get_cache_root()}")
  114. print(f"text_embedding 缓存: {get_cache_dir('text_embedding')}")
  115. print()
  116. # 示例 4: 通过环境变量设置
  117. print("示例 4: 通过环境变量设置")
  118. Config.reset() # 重置配置
  119. os.environ["CACHE_ROOT"] = "/Users/semsevens/Desktop/workspace/daily/1113/how_1120_v3/cache"
  120. print(f"缓存根目录: {get_cache_root()}")
  121. print(f"text_embedding 缓存: {get_cache_dir('text_embedding')}")
  122. print()
  123. # 示例 5: 数据目录配置(在缓存根目录下)
  124. print("示例 5: 数据目录配置(在缓存根目录下)")
  125. Config.reset() # 重置配置
  126. print(f"缓存根目录: {get_cache_root()}")
  127. print(f"data 目录: {get_data_dir()}")
  128. print(f"search 数据: {get_data_dir('search')}")
  129. print(f"detail 数据: {get_data_dir('detail')}")
  130. print()
  131. # 示例 6: 设置缓存根目录后,data 也会跟着变
  132. print("示例 6: 设置缓存根目录后,data 也会跟着变")
  133. set_cache_root("/custom/cache")
  134. print(f"缓存根目录: {get_cache_root()}")
  135. print(f"data 目录: {get_data_dir()}")
  136. print(f"search 数据: {get_data_dir('search')}")
  137. print()
  138. print("=" * 60)
  139. print("使用方法:")
  140. print("-" * 60)
  141. print("缓存根目录:")
  142. print(" 1. 默认使用 'cache' 目录")
  143. print(" 2. 通过代码设置: set_cache_root('/path/to/cache')")
  144. print(" 3. 通过环境变量: export CACHE_ROOT=/path/to/cache")
  145. print()
  146. print("目录结构:")
  147. print(" cache/")
  148. print(" ├── text_embedding/ # 向量相似度缓存")
  149. print(" ├── semantic_similarity/ # 语义相似度缓存")
  150. print(" └── data/ # 数据缓存(原 data 目录)")
  151. print(" ├── search/ # 搜索数据")
  152. print(" ├── detail/ # 详情数据")
  153. print(" └── analysis/ # 分析结果")
  154. print("=" * 60)