123 lines
4.1 KiB
Python
123 lines
4.1 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
脚本输出捕获工具
|
||
用于捕获脚本执行过程中的标准输出和错误输出，并通过WebSocket推送给订阅的客户端
|
||
"""
|
||
|
||
import asyncio
|
||
import sys
|
||
from io import StringIO
|
||
from contextlib import contextmanager
|
||
from .script_websocket_service import get_websocket_manager
|
||
from utils.logger import get_logger
|
||
|
||
logger = get_logger("services.script_output_capture")
|
||
|
||
|
||
class ScriptOutputCapture:
|
||
"""脚本输出捕获类"""
|
||
|
||
def __init__(self, script_id: str, level: str = "INFO", original_stream=None):
|
||
self.script_id = script_id
|
||
self.level = level
|
||
self.websocket_manager = get_websocket_manager()
|
||
self.buffer = StringIO()
|
||
self.original_stream = original_stream # 保留原始输出流
|
||
|
||
def write(self, text):
|
||
"""写入文本，同时推送到WebSocket和原始输出流"""
|
||
# 1. 写入原始输出流（如控制台）
|
||
if self.original_stream:
|
||
try:
|
||
self.original_stream.write(text)
|
||
self.original_stream.flush()
|
||
except Exception as e:
|
||
logger.debug(f"写入原始输出流失败: {e}")
|
||
|
||
# 2. 推送到WebSocket（仅对非空内容）
|
||
if text and text.strip():
|
||
try:
|
||
# 尝试在当前事件循环中创建任务
|
||
asyncio.create_task(self._push_output(text.strip()))
|
||
except RuntimeError:
|
||
# 如果没有运行中的事件循环，记录日志但不报错
|
||
logger.debug(f"无法推送输出到WebSocket，没有运行中的事件循环: [{self.script_id}] {text.strip()}")
|
||
|
||
# 3. 写入缓冲区保留原有行为
|
||
self.buffer.write(text)
|
||
|
||
def flush(self):
|
||
"""刷新缓冲区"""
|
||
self.buffer.flush()
|
||
if self.original_stream:
|
||
try:
|
||
self.original_stream.flush()
|
||
except Exception:
|
||
pass
|
||
|
||
async def _push_output(self, message: str):
|
||
"""推送输出到WebSocket"""
|
||
try:
|
||
await self.websocket_manager.broadcast_script_log(
|
||
self.script_id, self.level, message
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"推送脚本输出失败 [{self.script_id}]: {e}")
|
||
|
||
|
||
@contextmanager
|
||
def capture_script_output(script_id: str):
|
||
"""
|
||
上下文管理器，用于捕获脚本执行过程中的标准输出和错误输出
|
||
|
||
使用示例:
|
||
with capture_script_output("script_123"):
|
||
print("这条消息会被捕获并推送到WebSocket，同时显示在控制台")
|
||
# 执行脚本代码
|
||
"""
|
||
# 备份原始输出流
|
||
original_stdout = sys.stdout
|
||
original_stderr = sys.stderr
|
||
|
||
# 创建捕获器，传入原始输出流以保持双重输出
|
||
stdout_capture = ScriptOutputCapture(script_id, "INFO", original_stdout)
|
||
stderr_capture = ScriptOutputCapture(script_id, "ERROR", original_stderr)
|
||
|
||
try:
|
||
# 替换标准输出流
|
||
sys.stdout = stdout_capture
|
||
sys.stderr = stderr_capture
|
||
|
||
yield {
|
||
'stdout_capture': stdout_capture,
|
||
'stderr_capture': stderr_capture
|
||
}
|
||
|
||
finally:
|
||
# 恢复原始输出流
|
||
sys.stdout = original_stdout
|
||
sys.stderr = original_stderr
|
||
|
||
|
||
async def push_script_error(script_id: str, error_message: str):
|
||
"""推送脚本错误信息到WebSocket"""
|
||
try:
|
||
websocket_manager = get_websocket_manager()
|
||
await websocket_manager.broadcast_script_log(
|
||
script_id, "ERROR", error_message
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"推送脚本错误信息失败 [{script_id}]: {e}")
|
||
|
||
|
||
async def push_script_info(script_id: str, info_message: str):
|
||
"""推送脚本信息到WebSocket"""
|
||
try:
|
||
websocket_manager = get_websocket_manager()
|
||
await websocket_manager.broadcast_script_log(
|
||
script_id, "INFO", info_message
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"推送脚本信息失败 [{script_id}]: {e}") |