1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
| import asyncio
import json
import base64
import os
from PIL import Image
import io
class CDPScreenshotPDF:
"""CDP 截图与 PDF 工具"""
def __init__(self, ws, session_id):
self.ws = ws
self.session_id = session_id
self._cmd_id = 0
async def _cmd(self, method, params=None):
self._cmd_id += 1
msg = {
"sessionId": self.session_id,
"id": self._cmd_id,
"method": method,
"params": params or {}
}
await self.ws.send(json.dumps(msg))
async for resp in self.ws:
data = json.loads(resp)
if data.get("id") == self._cmd_id:
return data.get("result", {})
async def screenshot(self, path, format="png", quality=None):
params = {"format": format, "captureBeyondViewport": True}
if quality and format == "jpeg":
params["quality"] = quality
result = await self._cmd("Page.captureScreenshot", params)
data = base64.b64decode(result["data"])
with open(path, "wb") as f:
f.write(data)
return path
async def fullpage_screenshot(self, path):
result = await self._cmd("Runtime.evaluate", {
"expression": "JSON.stringify({w: document.documentElement.scrollWidth, h: document.documentElement.scrollHeight})",
"returnByValue": True
})
size = json.loads(result["result"]["value"])
await self._cmd("Emulation.setDeviceMetricsOverride", {
"width": size["w"], "height": size["h"],
"deviceScaleFactor": 1, "mobile": False
})
await asyncio.sleep(0.5)
result = await self._cmd("Page.captureScreenshot", {
"format": "png", "captureBeyondViewport": True
})
await self._cmd("Emulation.clearDeviceMetricsOverride")
data = base64.b64decode(result["data"])
with open(path, "wb") as f:
f.write(data)
return path
async def element_screenshot(self, selector, path):
result = await self._cmd("Runtime.evaluate", {
"expression": f"""
JSON.stringify(
(el => el ? {{x: el.getBoundingClientRect().x, y: el.getBoundingClientRect().y, w: el.offsetWidth, h: el.offsetHeight}} : null)
(document.querySelector('{selector}'))
)
""",
"returnByValue": True
})
rect = json.loads(result["result"]["value"])
if not rect:
return None
shot = await self._cmd("Page.captureScreenshot", {"format": "png"})
img = Image.open(io.BytesIO(base64.b64decode(shot["data"])))
cropped = img.crop((rect["x"], rect["y"], rect["x"] + rect["w"], rect["y"] + rect["h"]))
cropped.save(path)
return path
async def pdf(self, path, **kwargs):
params = {"printBackground": True, **kwargs}
result = await self._cmd("Page.printToPDF", params)
data = base64.b64decode(result["data"])
with open(path, "wb") as f:
f.write(data)
return path
|
