维护一个简单的基于令牌的文件下载服务,支持超时和懒清除。
| 5 | |
| 6 | |
| 7 | class FileTokenService: |
| 8 | """维护一个简单的基于令牌的文件下载服务,支持超时和懒清除。""" |
| 9 | |
| 10 | def __init__(self, default_timeout: float = 300) -> None: |
| 11 | self.lock = asyncio.Lock() |
| 12 | self.staged_files = {} # token: (file_path, expire_time) |
| 13 | self.default_timeout = default_timeout |
| 14 | |
| 15 | async def _cleanup_expired_tokens(self) -> None: |
| 16 | """清理过期的令牌""" |
| 17 | now = time.time() |
| 18 | expired_tokens = [ |
| 19 | token for token, (_, expire) in self.staged_files.items() if expire < now |
| 20 | ] |
| 21 | for token in expired_tokens: |
| 22 | self.staged_files.pop(token, None) |
| 23 | |
| 24 | async def check_token_expired(self, file_token: str) -> bool: |
| 25 | async with self.lock: |
| 26 | await self._cleanup_expired_tokens() |
| 27 | return file_token not in self.staged_files |
| 28 | |
| 29 | async def register_file(self, file_path: str, timeout: float | None = None) -> str: |
| 30 | """向令牌服务注册一个文件。 |
| 31 | |
| 32 | Args: |
| 33 | file_path(str): 文件路径 |
| 34 | timeout(float): 超时时间,单位秒(可选) |
| 35 | |
| 36 | Returns: |
| 37 | str: 一个单次令牌 |
| 38 | |
| 39 | Raises: |
| 40 | FileNotFoundError: 当路径不存在时抛出 |
| 41 | |
| 42 | """ |
| 43 | try: |
| 44 | from astrbot.core.utils.media_utils import file_uri_to_path, is_file_uri |
| 45 | |
| 46 | local_path = ( |
| 47 | file_uri_to_path(file_path) if is_file_uri(file_path) else file_path |
| 48 | ) |
| 49 | except Exception: |
| 50 | # Fall back to the original path if URL parsing fails. |
| 51 | local_path = file_path |
| 52 | |
| 53 | async with self.lock: |
| 54 | await self._cleanup_expired_tokens() |
| 55 | |
| 56 | if not os.path.exists(local_path): |
| 57 | raise FileNotFoundError( |
| 58 | f"文件不存在: {local_path} (原始输入: {file_path})", |
| 59 | ) |
| 60 | |
| 61 | file_token = str(uuid.uuid4()) |
| 62 | expire_time = time.time() + ( |
| 63 | timeout if timeout is not None else self.default_timeout |
| 64 | ) |
no outgoing calls