first commit
This commit is contained in:
207
utils/screenshot.py
Normal file
207
utils/screenshot.py
Normal file
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
截图工具模块
|
||||
|
||||
提供多种截图方式：
|
||||
1. rpaframework 内置截图（浏览器/桌面）
|
||||
2. pyautogui 全屏/区域截图
|
||||
3. pillow 窗口截图
|
||||
"""
|
||||
import os
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from RPA.Browser import Playwright
|
||||
from config import OUTPUT_DIR
|
||||
from utils import get_logger, now_str
|
||||
|
||||
log = get_logger("screenshot")
|
||||
|
||||
try:
|
||||
import pyautogui
|
||||
import pygetwindow as gw
|
||||
except ImportError:
|
||||
pyautogui = None
|
||||
gw = None
|
||||
|
||||
|
||||
def take_screenshot(
|
||||
name: Optional[str] = None,
|
||||
full_page: bool = False,
|
||||
region: Optional[Tuple[int, int, int, int]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
截图工具
|
||||
|
||||
Args:
|
||||
name: 截图文件名（不含扩展名），默认为时间戳
|
||||
full_page: 是否截取整个页面（仅浏览器有效）
|
||||
region: 区域截图 (left, top, width, height)，None为全屏
|
||||
|
||||
Returns:
|
||||
截图文件完整路径
|
||||
"""
|
||||
if name is None:
|
||||
name = now_str()
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
filepath = os.path.join(OUTPUT_DIR, f"{name}.png")
|
||||
|
||||
# 优先尝试 rpaframework 浏览器截图
|
||||
try:
|
||||
browser = Playwright()
|
||||
if browser.get_browser_count() > 0:
|
||||
# 如果有浏览器实例，则截取当前页面
|
||||
if full_page:
|
||||
browser.take_screenshot(filepath, selector="body")
|
||||
else:
|
||||
browser.take_screenshot(filepath)
|
||||
log.info("已使用 rpaframework 截取浏览器页面: %s", filepath)
|
||||
return filepath
|
||||
except:
|
||||
# 没有浏览器实例，继续尝试其他方式
|
||||
pass
|
||||
|
||||
# 尝试 pyautogui 截图
|
||||
if pyautogui:
|
||||
try:
|
||||
if region:
|
||||
# 区域截图
|
||||
screenshot = pyautogui.screenshot(region=region)
|
||||
else:
|
||||
# 全屏截图
|
||||
screenshot = pyautogui.screenshot()
|
||||
|
||||
screenshot.save(filepath)
|
||||
log.info("已使用 pyautogui 截取屏幕: %s", filepath)
|
||||
return filepath
|
||||
except Exception as e:
|
||||
log.warning("pyautogui 截图失败: %s", e)
|
||||
|
||||
# 尝试截取活动窗口
|
||||
if gw:
|
||||
try:
|
||||
active_window = gw.getActiveWindow()
|
||||
if active_window:
|
||||
bbox = (
|
||||
active_window.left,
|
||||
active_window.top,
|
||||
active_window.width,
|
||||
active_window.height,
|
||||
)
|
||||
screenshot = pyautogui.screenshot(region=bbox)
|
||||
filepath = os.path.join(OUTPUT_DIR, f"{name}_window.png")
|
||||
screenshot.save(filepath)
|
||||
log.info("已截取活动窗口: %s", filepath)
|
||||
return filepath
|
||||
except Exception as e:
|
||||
log.warning("窗口截图失败: %s", e)
|
||||
|
||||
# 最后尝试 Pillow 截图（需要安装 Pillow）
|
||||
try:
|
||||
from PIL import ImageGrab
|
||||
|
||||
if region:
|
||||
bbox = (region[0], region[1], region[0] + region[2], region[1] + region[3])
|
||||
img = ImageGrab.grab(bbox=bbox)
|
||||
else:
|
||||
img = ImageGrab.grab()
|
||||
|
||||
img.save(filepath)
|
||||
log.info("已使用 Pillow 截取屏幕: %s", filepath)
|
||||
return filepath
|
||||
except ImportError:
|
||||
log.error("未安装 Pillow，无法进行屏幕截图")
|
||||
raise
|
||||
except Exception as e:
|
||||
log.error("所有截图方式均失败: %s", e)
|
||||
raise
|
||||
|
||||
|
||||
def take_element_screenshot(selector: str, name: Optional[str] = None) -> str:
|
||||
"""
|
||||
截取浏览器中特定元素
|
||||
|
||||
Args:
|
||||
selector: CSS 选择器或 XPath
|
||||
name: 截图文件名
|
||||
|
||||
Returns:
|
||||
截图文件完整路径
|
||||
"""
|
||||
if name is None:
|
||||
name = f"element_{now_str()}"
|
||||
|
||||
filepath = os.path.join(OUTPUT_DIR, f"{name}.png")
|
||||
|
||||
try:
|
||||
browser = Playwright()
|
||||
if browser.get_browser_count() == 0:
|
||||
raise RuntimeError("没有活动的浏览器实例")
|
||||
|
||||
browser.take_screenshot(filepath, selector=selector)
|
||||
log.info("已截取元素 '%s': %s", selector, filepath)
|
||||
return filepath
|
||||
except Exception as e:
|
||||
log.error("元素截图失败: %s", e)
|
||||
raise
|
||||
|
||||
|
||||
def take_window_screenshot(window_title: str, name: Optional[str] = None) -> str:
|
||||
"""
|
||||
截取指定窗口
|
||||
|
||||
Args:
|
||||
window_title: 窗口标题（模糊匹配）
|
||||
name: 截图文件名
|
||||
|
||||
Returns:
|
||||
截图文件完整路径
|
||||
"""
|
||||
if name is None:
|
||||
name = f"window_{now_str()}"
|
||||
|
||||
if not gw:
|
||||
raise RuntimeError("未安装 pygetwindow，无法截取窗口")
|
||||
|
||||
try:
|
||||
# 查找匹配的窗口
|
||||
windows = gw.getWindowsWithTitle(window_title)
|
||||
if not windows:
|
||||
raise ValueError(f"未找到标题包含 '{window_title}' 的窗口")
|
||||
|
||||
# 使用第一个匹配的窗口
|
||||
target_window = windows[0]
|
||||
bbox = (
|
||||
target_window.left,
|
||||
target_window.top,
|
||||
target_window.width,
|
||||
target_window.height,
|
||||
)
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
filepath = os.path.join(OUTPUT_DIR, f"{name}.png")
|
||||
|
||||
# 截取窗口
|
||||
screenshot = pyautogui.screenshot(region=bbox)
|
||||
screenshot.save(filepath)
|
||||
|
||||
log.info("已截取窗口 '%s': %s", window_title, filepath)
|
||||
return filepath
|
||||
except Exception as e:
|
||||
log.error("窗口截图失败: %s", e)
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 示例：全屏截图
|
||||
take_screenshot("full_screen_demo")
|
||||
|
||||
# 示例：区域截图 (left, top, width, height)
|
||||
# take_screenshot("region_demo", region=(100, 100, 500, 400))
|
||||
|
||||
# 示例：窗口截图
|
||||
# try:
|
||||
# take_window_screenshot("微信", "wechat_window_demo")
|
||||
# except:
|
||||
# print("未找到微信窗口，跳过窗口截图示例")
|
||||
Reference in New Issue
Block a user