mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
refactor(utils): improve type annotations and code quality
command_parser.py: - Add explicit type annotations for tokens list and return type config_number.py: - Handle float to int conversion explicitly - Simplify type checking logic file_extract.py: - Add type annotations and improve file extraction handling io.py: - Add type annotations for I/O operations log_pipe.py: - Implement TextIO interface for better compatibility - Add comprehensive type annotations - Add callback support and identifier logging session_waiter.py: - Clean up type annotations
This commit is contained in:
@@ -3,7 +3,7 @@ import re
|
||||
|
||||
class CommandTokens:
|
||||
def __init__(self) -> None:
|
||||
self.tokens = []
|
||||
self.tokens: list[str] = []
|
||||
self.len = 0
|
||||
|
||||
def get(self, idx: int) -> str | None:
|
||||
@@ -13,7 +13,7 @@ class CommandTokens:
|
||||
|
||||
|
||||
class CommandParserMixin:
|
||||
def parse_commands(self, message: str):
|
||||
def parse_commands(self, message: str) -> CommandTokens:
|
||||
cmd_tokens = CommandTokens()
|
||||
cmd_tokens.tokens = re.split(r"\s+", message)
|
||||
cmd_tokens.len = len(cmd_tokens.tokens)
|
||||
|
||||
@@ -36,19 +36,18 @@ def coerce_int_config(
|
||||
default,
|
||||
)
|
||||
parsed = default
|
||||
elif isinstance(value, float):
|
||||
parsed = int(value)
|
||||
else:
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
if warn:
|
||||
logger.warning(
|
||||
"%s %s has unsupported type %s. Fallback to %s.",
|
||||
source,
|
||||
label,
|
||||
type(value).__name__,
|
||||
default,
|
||||
)
|
||||
parsed = default
|
||||
if warn:
|
||||
logger.warning(
|
||||
"%s %s has unsupported type %s. Fallback to %s.",
|
||||
source,
|
||||
label,
|
||||
type(value).__name__,
|
||||
default,
|
||||
)
|
||||
parsed = default
|
||||
|
||||
if min_value is not None and parsed < min_value:
|
||||
if warn:
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from pathlib import Path
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
import anyio
|
||||
import httpx
|
||||
|
||||
|
||||
async def extract_file_moonshotai(file_path: str, api_key: str) -> str:
|
||||
@@ -12,12 +13,36 @@ async def extract_file_moonshotai(file_path: str, api_key: str) -> str:
|
||||
Returns:
|
||||
The text extracted from the file
|
||||
"""
|
||||
client = AsyncOpenAI(
|
||||
api_key=api_key,
|
||||
base_url="https://api.moonshot.cn/v1",
|
||||
)
|
||||
file_object = await client.files.create(
|
||||
file=Path(file_path),
|
||||
purpose="file-extract",
|
||||
)
|
||||
return (await client.files.content(file_id=file_object.id)).text
|
||||
base_url = "https://api.moonshot.cn/v1"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
source_path = Path(file_path)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
base_url=base_url,
|
||||
headers=headers,
|
||||
follow_redirects=True,
|
||||
timeout=60.0,
|
||||
) as client:
|
||||
source_bytes = await anyio.Path(source_path).read_bytes()
|
||||
upload_response = await client.post(
|
||||
"/files",
|
||||
data={"purpose": "file-extract"},
|
||||
files={
|
||||
"file": (
|
||||
source_path.name,
|
||||
source_bytes,
|
||||
"application/octet-stream",
|
||||
)
|
||||
},
|
||||
)
|
||||
upload_response.raise_for_status()
|
||||
uploaded_file = upload_response.json()
|
||||
file_id = uploaded_file.get("id")
|
||||
if not isinstance(file_id, str) or not file_id:
|
||||
raise ValueError("Moonshot file upload did not return a valid file id")
|
||||
|
||||
content_response = await client.get(f"/files/{file_id}/content")
|
||||
content_response.raise_for_status()
|
||||
return content_response.text
|
||||
|
||||
@@ -145,7 +145,9 @@ async def download_file(url: str, path: str, show_progress: bool = False) -> Non
|
||||
trust_env=True,
|
||||
connector=connector,
|
||||
) as session:
|
||||
async with session.get(url, timeout=1800) as resp:
|
||||
async with session.get(
|
||||
url, timeout=aiohttp.ClientTimeout(total=1800)
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
raise Exception(f"下载文件失败: {resp.status}")
|
||||
total_size = int(resp.headers.get("content-length", 0))
|
||||
@@ -187,7 +189,11 @@ async def download_file(url: str, path: str, show_progress: bool = False) -> Non
|
||||
ssl_context.check_hostname = False
|
||||
ssl_context.verify_mode = ssl.CERT_NONE
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, ssl=ssl_context, timeout=120) as resp:
|
||||
async with session.get(
|
||||
url,
|
||||
ssl=ssl_context,
|
||||
timeout=aiohttp.ClientTimeout(total=120),
|
||||
) as resp:
|
||||
total_size = int(resp.headers.get("content-length", 0))
|
||||
downloaded_size = 0
|
||||
start_time = time.time()
|
||||
@@ -247,7 +253,7 @@ async def get_public_ip_address() -> list[IPv4Address | IPv6Address]:
|
||||
|
||||
async def fetch(session: aiohttp.ClientSession, url: str):
|
||||
try:
|
||||
async with session.get(url, timeout=3) as resp:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
|
||||
if resp.status == 200:
|
||||
raw_ip = (await resp.text()).strip()
|
||||
ip = ip_address(raw_ip)
|
||||
|
||||
@@ -1,36 +1,96 @@
|
||||
import io
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Callable, Iterable
|
||||
from logging import Logger
|
||||
from types import TracebackType
|
||||
from typing import Self
|
||||
|
||||
|
||||
class LogPipe(threading.Thread):
|
||||
class LogPipe(threading.Thread, io.TextIOBase):
|
||||
"""A pipe wrapper that routes written content to a logger.
|
||||
|
||||
Implements TextIO interface for compatibility with code expecting
|
||||
a text stream, while also logging all written content.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
level,
|
||||
level: int,
|
||||
logger: Logger,
|
||||
identifier=None,
|
||||
callback=None,
|
||||
identifier: str | None = None,
|
||||
callback: Callable[[str], None] | None = None,
|
||||
) -> None:
|
||||
threading.Thread.__init__(self)
|
||||
self.daemon = True
|
||||
self.level = level
|
||||
self._logger = logger
|
||||
self._identifier = identifier
|
||||
self._callback = callback
|
||||
self._closed = False
|
||||
self.fd_read, self.fd_write = os.pipe()
|
||||
self.identifier = identifier
|
||||
self.logger = logger
|
||||
self.callback = callback
|
||||
self.reader = os.fdopen(self.fd_read)
|
||||
self._reader = os.fdopen(self.fd_read, "r")
|
||||
self.mode = "w"
|
||||
self.name = f"<LogPipe-{self._identifier}>"
|
||||
self.start()
|
||||
|
||||
def fileno(self):
|
||||
def fileno(self) -> int:
|
||||
return self.fd_write
|
||||
|
||||
def run(self) -> None:
|
||||
for line in iter(self.reader.readline, ""):
|
||||
if self.callback:
|
||||
self.callback(line.strip())
|
||||
self.logger.log(self.level, f"[{self.identifier}] {line.strip()}")
|
||||
def write(self, s: str) -> int:
|
||||
"""Write string to pipe - content will be logged."""
|
||||
if self._closed:
|
||||
raise ValueError("I/O operation on closed file")
|
||||
self._logger.log(self.level, f"[{self._identifier}] {s.rstrip()}")
|
||||
if self._callback:
|
||||
self._callback(s.strip())
|
||||
return len(s)
|
||||
|
||||
self.reader.close()
|
||||
def flush(self) -> None:
|
||||
"""No-op for compatibility - log writes are immediate."""
|
||||
pass
|
||||
|
||||
def close(self) -> None:
|
||||
os.close(self.fd_write)
|
||||
"""Close the write end of the pipe."""
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
os.close(self.fd_write)
|
||||
|
||||
def isatty(self) -> bool:
|
||||
return False
|
||||
|
||||
def writable(self) -> bool:
|
||||
return not self._closed
|
||||
|
||||
def readable(self) -> bool:
|
||||
return False
|
||||
|
||||
def seekable(self) -> bool:
|
||||
return False
|
||||
|
||||
def writelines(self, lines: Iterable[str]) -> None:
|
||||
for line in lines:
|
||||
self.write(line)
|
||||
|
||||
def run(self) -> None:
|
||||
"""Read from pipe and log each line."""
|
||||
for line in iter(self._reader.readline, ""):
|
||||
if self._closed:
|
||||
break
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
self._logger.log(self.level, f"[{self._identifier}] {stripped}")
|
||||
if self._callback:
|
||||
self._callback(stripped)
|
||||
self._reader.close()
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.close()
|
||||
|
||||
@@ -19,8 +19,8 @@ class SessionController:
|
||||
"""控制一个 Session 是否已经结束"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.tasks = set()
|
||||
self.future = asyncio.Future()
|
||||
self.tasks: set[asyncio.Task[None]] = set()
|
||||
self.future: asyncio.Future[None] = asyncio.Future()
|
||||
self.current_event: asyncio.Event | None = None
|
||||
"""当前正在等待的所用的异步事件"""
|
||||
self.ts: float | None = None
|
||||
|
||||
Reference in New Issue
Block a user