|
1 | 1 | from __future__ import annotations |
2 | 2 | import asyncio |
| 3 | +import base64 |
| 4 | +import datetime |
3 | 5 | import logging |
4 | 6 | import pathlib |
| 7 | +import urllib.parse |
5 | 8 | import warnings |
6 | 9 | from typing import Dict, List, Union, Optional, Tuple |
7 | 10 | from . import browser as cdp_browser |
@@ -1133,9 +1136,6 @@ async def save_screenshot( |
1133 | 1136 | :return: The path/filename of the saved screenshot. |
1134 | 1137 | :rtype: str |
1135 | 1138 | """ |
1136 | | - import urllib.parse |
1137 | | - import datetime |
1138 | | - |
1139 | 1139 | await self.sleep() # Update the target's URL |
1140 | 1140 | path = None |
1141 | 1141 | if format.lower() in ["jpg", "jpeg"]: |
@@ -1166,8 +1166,40 @@ async def save_screenshot( |
1166 | 1166 | "Most possible cause is the page " |
1167 | 1167 | "has not finished loading yet." |
1168 | 1168 | ) |
1169 | | - import base64 |
| 1169 | + data_bytes = base64.b64decode(data) |
| 1170 | + if not path: |
| 1171 | + raise RuntimeError("Invalid filename or path: '%s'" % filename) |
| 1172 | + path.write_bytes(data_bytes) |
| 1173 | + return str(path) |
1170 | 1174 |
|
| 1175 | + async def print_to_pdf( |
| 1176 | + self, |
| 1177 | + filename: Optional[PathLike] = "auto", |
| 1178 | + ) -> str: |
| 1179 | + """ |
| 1180 | + Saves a webpage as a PDF. |
| 1181 | + :param filename: uses this as the save path |
| 1182 | + :type filename: PathLike |
| 1183 | + :return: The path/filename of the saved screenshot. |
| 1184 | + :rtype: str |
| 1185 | + """ |
| 1186 | + await self.sleep() # Update the target's URL |
| 1187 | + path = None |
| 1188 | + ext = ".pdf" |
| 1189 | + if not filename or filename == "auto": |
| 1190 | + parsed = urllib.parse.urlparse(self.target.url) |
| 1191 | + parts = parsed.path.split("/") |
| 1192 | + last_part = parts[-1] |
| 1193 | + last_part = last_part.rsplit("?", 1)[0] |
| 1194 | + dt_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
| 1195 | + candidate = f"{parsed.hostname}__{last_part}_{dt_str}" |
| 1196 | + path = pathlib.Path(candidate + ext) # noqa |
| 1197 | + else: |
| 1198 | + path = pathlib.Path(filename) |
| 1199 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 1200 | + data, _ = await self.send(cdp.page.print_to_pdf()) |
| 1201 | + if not data: |
| 1202 | + raise ProtocolException("Could not save PDF.") |
1171 | 1203 | data_bytes = base64.b64decode(data) |
1172 | 1204 | if not path: |
1173 | 1205 | raise RuntimeError("Invalid filename or path: '%s'" % filename) |
|
0 commit comments