208 lines
6.1 KiB
Python
208 lines
6.1 KiB
Python
"""
|
||
截图工具模块
|
||
|
||
提供多种截图方式：
|
||
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("未找到微信窗口，跳过窗口截图示例")
|