feat: 添加了支持选择元素截图和等待资源加载完成之后再截图的功能#14
Open
kyuphoenix wants to merge 2 commits into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces element-specific screenshots using CSS selectors (with fallback options and timeouts) and resource waiting (waiting for network idle, image decoding, and font loading before taking a screenshot). The review feedback suggests several robustness improvements: wrapping query_selector calls in try-except blocks to handle invalid CSS selectors gracefully, avoiding object destructuring in the JavaScript evaluation string to prevent Playwright parsing issues, and verifying that the result of page.evaluate is a dictionary before calling .get() to avoid potential AttributeError exceptions.
Comment on lines
+298
to
+336
| async def _resolve_screenshot_element(self, page, screenshot_options: ScreenshotOptions): | ||
| """Resolve an element target for screenshot when selector options are provided.""" | ||
|
|
||
| if not screenshot_options.selector: | ||
| return None | ||
|
|
||
| element = None | ||
| if screenshot_options.selector_timeout is not None: | ||
| try: | ||
| element = await page.wait_for_selector( | ||
| screenshot_options.selector, | ||
| timeout=screenshot_options.selector_timeout, | ||
| ) | ||
| except Exception as e: | ||
| logger.debug( | ||
| f"html2pic: wait for selector '{screenshot_options.selector}' failed: {e}" | ||
| ) | ||
| else: | ||
| element = await page.query_selector(screenshot_options.selector) | ||
|
|
||
| if element is not None: | ||
| logger.info( | ||
| f"html2pic: screenshot element matched selector '{screenshot_options.selector}'" | ||
| ) | ||
| return element | ||
|
|
||
| if screenshot_options.fallback_selector: | ||
| fallback = await page.query_selector(screenshot_options.fallback_selector) | ||
| if fallback is not None: | ||
| logger.info( | ||
| f"html2pic: selector '{screenshot_options.selector}' not found, " | ||
| f"screenshot fallback selector '{screenshot_options.fallback_selector}'" | ||
| ) | ||
| return fallback | ||
|
|
||
| logger.warning( | ||
| f"html2pic: selector '{screenshot_options.selector}' not found, fallback to page screenshot" | ||
| ) | ||
| return None |
There was a problem hiding this comment.
如果用户传入了非法的 CSS 选择器（例如语法错误的选择器），page.query_selector 会抛出异常。在当前实现中，如果 selector_timeout 为 None，或者在解析 fallback_selector 时，这些 query_selector 调用没有被 try...except 保护，会导致整个渲染请求崩溃并返回 500 错误。建议对所有的 query_selector 调用进行异常捕获，确保在选择器非法时能够优雅地记录警告并回退到整页截图。
async def _resolve_screenshot_element(self, page, screenshot_options: ScreenshotOptions):
"""Resolve an element target for screenshot when selector options are provided."""
if not screenshot_options.selector:
return None
element = None
if screenshot_options.selector_timeout is not None:
try:
element = await page.wait_for_selector(
screenshot_options.selector,
timeout=screenshot_options.selector_timeout,
)
except Exception as e:
logger.debug(
f"html2pic: wait for selector '{screenshot_options.selector}' failed: {e}"
)
else:
try:
element = await page.query_selector(screenshot_options.selector)
except Exception as e:
logger.warning(
f"html2pic: query selector '{screenshot_options.selector}' failed: {e}"
)
if element is not None:
logger.info(
f"html2pic: screenshot element matched selector '{screenshot_options.selector}'"
)
return element
if screenshot_options.fallback_selector:
try:
fallback = await page.query_selector(screenshot_options.fallback_selector)
if fallback is not None:
logger.info(
f"html2pic: selector '{screenshot_options.selector}' not found, "
f"screenshot fallback selector '{screenshot_options.fallback_selector}'"
)
return fallback
except Exception as e:
logger.warning(
f"html2pic: query fallback selector '{screenshot_options.fallback_selector}' failed: {e}"
)
logger.warning(
f"html2pic: selector '{screenshot_options.selector}' not found, fallback to page screenshot"
)
return None
Comment on lines
+365
to
+408
| result = await page.evaluate( | ||
| """ | ||
| async ({ timeout }) => { | ||
| const images = Array.from(document.images || []); | ||
| const waitImage = (img) => { | ||
| if (img.complete) { | ||
| return Promise.resolve(); | ||
| } | ||
| if (typeof img.decode === "function") { | ||
| return img.decode().catch(() => undefined); | ||
| } | ||
| return new Promise((resolve) => { | ||
| img.addEventListener("load", resolve, { once: true }); | ||
| img.addEventListener("error", resolve, { once: true }); | ||
| }); | ||
| }; | ||
| const waitFonts = () => { | ||
| if (document.fonts && document.fonts.ready) { | ||
| return document.fonts.ready.catch(() => undefined); | ||
| } | ||
| return Promise.resolve(); | ||
| }; | ||
| let timedOut = false; | ||
| const allResources = Promise.all([ | ||
| ...images.map(waitImage), | ||
| waitFonts(), | ||
| ]); | ||
| const timeoutPromise = new Promise((resolve) => { | ||
| setTimeout(() => { | ||
| timedOut = true; | ||
| resolve(); | ||
| }, timeout); | ||
| }); | ||
| await Promise.race([allResources, timeoutPromise]); | ||
| return { | ||
| imageCount: images.length, | ||
| completeCount: images.filter((img) => img.complete).length, | ||
| brokenCount: images.filter((img) => img.complete && img.naturalWidth === 0).length, | ||
| timedOut, | ||
| }; | ||
| } | ||
| """, | ||
| {"timeout": resource_timeout}, | ||
| ) |
There was a problem hiding this comment.
在 Playwright Python 中，使用字符串传递 JS 函数时，如果函数参数中包含对象解构（如 async ({ timeout }) =>），可能会导致 Playwright 的函数检测启发式算法失效，从而无法正确执行该函数或引发解析错误。建议改用普通的参数列表（如 async (timeout) =>），并将 resource_timeout 直接作为参数传入。
result = await page.evaluate(
"""
async (timeout) => {
const images = Array.from(document.images || []);
const waitImage = (img) => {
if (img.complete) {
return Promise.resolve();
}
if (typeof img.decode === "function") {
return img.decode().catch(() => undefined);
}
return new Promise((resolve) => {
img.addEventListener("load", resolve, { once: true });
img.addEventListener("error", resolve, { once: true });
});
};
const waitFonts = () => {
if (document.fonts && document.fonts.ready) {
return document.fonts.ready.catch(() => undefined);
}
return Promise.resolve();
};
let timedOut = false;
const allResources = Promise.all([
...images.map(waitImage),
waitFonts(),
]);
const timeoutPromise = new Promise((resolve) => {
setTimeout(() => {
timedOut = true;
resolve();
}, timeout);
});
await Promise.race([allResources, timeoutPromise]);
return {
imageCount: images.length,
completeCount: images.filter((img) => img.complete).length,
brokenCount: images.filter((img) => img.complete && img.naturalWidth === 0).length,
timedOut,
};
}
""",
resource_timeout,
)
Comment on lines
+409
to
+426
| if result.get("timedOut"): | ||
| logger.warning( | ||
| "html2pic: wait for resources timed out after " | ||
| f"{timeout}ms, images={result.get('completeCount')}/" | ||
| f"{result.get('imageCount')}" | ||
| ) | ||
| elif result.get("brokenCount"): | ||
| logger.warning( | ||
| "html2pic: resources loaded with broken images, " | ||
| f"broken={result.get('brokenCount')}/" | ||
| f"{result.get('imageCount')}" | ||
| ) | ||
| else: | ||
| logger.info( | ||
| "html2pic: resources ready, " | ||
| f"images={result.get('completeCount')}/" | ||
| f"{result.get('imageCount')}" | ||
| ) |
There was a problem hiding this comment.
page.evaluate 的返回值 result 在某些异常情况下（例如 JS 执行未返回预期对象或返回了 null/undefined）可能为 None。如果直接调用 result.get("timedOut") 会抛出 AttributeError。虽然该代码块处于 try...except 中，但这会导致打印出不必要的错误日志，掩盖了真实的执行状态。建议在调用 .get() 之前先检查 result 是否为字典类型。
Suggested change
| if result.get("timedOut"): | |
| logger.warning( | |
| "html2pic: wait for resources timed out after " | |
| f"{timeout}ms, images={result.get('completeCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) | |
| elif result.get("brokenCount"): | |
| logger.warning( | |
| "html2pic: resources loaded with broken images, " | |
| f"broken={result.get('brokenCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) | |
| else: | |
| logger.info( | |
| "html2pic: resources ready, " | |
| f"images={result.get('completeCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) | |
| if not isinstance(result, dict): | |
| logger.warning("html2pic: wait for resources returned invalid result") | |
| elif result.get("timedOut"): | |
| logger.warning( | |
| "html2pic: wait for resources timed out after " | |
| f"{timeout}ms, images={result.get('completeCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) | |
| elif result.get("brokenCount"): | |
| logger.warning( | |
| "html2pic: resources loaded with broken images, " | |
| f"broken={result.get('brokenCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) | |
| else: | |
| logger.info( | |
| "html2pic: resources ready, " | |
| f"images={result.get('completeCount')}/" | |
| f"{result.get('imageCount')}" | |
| ) |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- In
_resolve_screenshot_element, thewait_for_selectorcall catches a broadExceptionwhich can hide non-timeout issues; consider catchingTimeoutError(or Playwright’s specific timeout exception) instead so real selector errors still surface. selector_timeoutandresource_timeoutare typed asfloatmilliseconds but are ultimately passed to Playwright as integer timeouts; it may be clearer and less error‑prone to standardize these as integer millisecond values or document/convert them explicitly at the boundary.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `_resolve_screenshot_element`, the `wait_for_selector` call catches a broad `Exception` which can hide non-timeout issues; consider catching `TimeoutError` (or Playwright’s specific timeout exception) instead so real selector errors still surface.
- `selector_timeout` and `resource_timeout` are typed as `float` milliseconds but are ultimately passed to Playwright as integer timeouts; it may be clearer and less error‑prone to standardize these as integer millisecond values or document/convert them explicitly at the boundary.
## Individual Comments
### Comment 1
<location path="src/render.py" line_range="305-311" />
<code_context>
+ return None
+
+ element = None
+ if screenshot_options.selector_timeout is not None:
+ try:
+ element = await page.wait_for_selector(
+ screenshot_options.selector,
+ timeout=screenshot_options.selector_timeout,
+ )
+ except Exception as e:
+ logger.debug(
+ f"html2pic: wait for selector '{screenshot_options.selector}' failed: {e}"
</code_context>
<issue_to_address>
**issue (bug_risk):** Narrow the exception handling around `wait_for_selector` to avoid masking non-timeout errors.
Catching `Exception` here turns all failures into a “selector not found” case, which can hide real problems (e.g., invalid selector, detached frame, other runtime errors) and only trigger the fallback full-page screenshot. Please catch only the timeout-related exception type (e.g., `TimeoutError` or the appropriate Playwright timeout error) and let other exceptions propagate so genuine bugs aren’t silently swallowed.
</issue_to_address>
### Comment 2
<location path="src/render.py" line_range="365-374" />
<code_context>
+ result = await page.evaluate(
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Align Playwright’s evaluation timeout with the custom `resource_timeout` to avoid unexpected early termination.
`resource_timeout` is only passed into the page context; Playwright’s own `page.evaluate` timeout remains at its default and may terminate the call before the in-page timeout triggers. Consider also passing `timeout=resource_timeout` (or `timeout=min(resource_timeout, screenshot_options.timeout or resource_timeout)`) to `page.evaluate` so both layers enforce the same max duration.
Suggested implementation:
```python
resource_timeout = remaining_timeout()
if resource_timeout <= 0:
logger.warning(f"html2pic: wait for resources timed out after {timeout}ms")
return
# Align Playwright's evaluation timeout with the in-page timeout
# If screenshot_options.timeout is set, clamp to the smaller value
eval_timeout = resource_timeout
if "screenshot_options" in locals() and getattr(screenshot_options, "timeout", None) is not None:
eval_timeout = min(resource_timeout, screenshot_options.timeout)
try:
result = await page.evaluate(
"""
async ({ timeout }) => {
const images = Array.from(document.images || []);
const waitImage = (img) => {
if (img.complete) {
return Promise.resolve();
}
if (typeof img.decode === "function") {
return img.decode().catch(() => undefined);
}
```
```python
result = await page.evaluate(
"""
async ({ timeout }) => {
const images = Array.from(document.images || []);
const waitImage = (img) => {
if (img.complete) {
return Promise.resolve();
}
if (typeof img.decode === "function") {
return img.decode().catch(() => undefined);
}
```
The snippet you provided stops inside the JavaScript function body, so I cannot see the full `page.evaluate` call. To fully implement the timeout alignment, you should also:
1. Ensure that when you pass the `timeout` into the page context, you use `eval_timeout` instead of `resource_timeout`, for example:
- Change the argument object from:
- `{"timeout": resource_timeout}`
- to:
- `{"timeout": eval_timeout}`
2. Add the Playwright timeout keyword argument to the `page.evaluate` call, so the final call shape is:
- `result = await page.evaluate(<script>, {"timeout": eval_timeout}, timeout=eval_timeout)`
Adjust for your actual argument formatting/indentation.
3. If `screenshot_options` is not available in this scope or its API is different (e.g., `screenshot_options.get("timeout")` instead of attribute access), update the `eval_timeout` calculation accordingly, preserving the intent:
- clamp Playwright’s timeout to `min(resource_timeout, configured_screenshot_timeout)` and fall back to `resource_timeout` when no explicit timeout is configured.
These small follow-up adjustments will ensure both the in-page timeout and Playwright’s own timeout enforce the same maximum duration.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
本次更新为 T2I 渲染服务新增了两组可选配置，旨在提升截图的灵活性与渲染稳定性。
修改后如果不传入对应配置项则保持功能与原先相同，避免因更新服务导致插件失效
Summary by Sourcery
Add configurable element-based screenshot targeting and optional resource-loading waits to the T2I HTML rendering flow while preserving existing defaults when unset.
New Features:
Enhancements:
Documentation: