Fix response encoding detection (#65)

elacuesta · web-flow · commit ef7542605829 · 2022-03-02T15:41:59.000-03:00
diff --git a/scrapy_playwright/handler.py b/scrapy_playwright/handler.py
@@ -22,8 +22,10 @@
 from scrapy.responsetypes import responsetypes
 from scrapy.utils.defer import deferred_from_coro
 from scrapy.utils.misc import load_object
+from scrapy.utils.python import to_unicode
 from scrapy.utils.reactor import verify_installed_reactor
 from twisted.internet.defer import Deferred, inlineCallbacks
+from w3lib.encoding import html_body_declared_encoding, http_content_type_encoding
 
 from scrapy_playwright.headers import use_scrapy_headers
 from scrapy_playwright.page import PageCoroutine
@@ -220,7 +222,7 @@ async def _download_request_with_page(self, request: Request, page: Page) -> Res
                 pc.result = await method(*pc.args, **pc.kwargs)
                 await page.wait_for_load_state(timeout=self.default_navigation_timeout)
 
-        body = (await page.content()).encode("utf8")
+        body_str = await page.content()
         request.meta["download_latency"] = time() - start_time
 
         if request.meta.get("playwright_include_page"):
@@ -231,6 +233,8 @@ async def _download_request_with_page(self, request: Request, page: Page) -> Res
 
         headers = Headers(response.headers)
         headers.pop("Content-Encoding", None)
+        encoding = _get_response_encoding(headers, body_str) or "utf-8"
+        body = body_str.encode(encoding)
         respcls = responsetypes.from_args(headers=headers, url=page.url, body=body)
         return respcls(
             url=page.url,
@@ -239,6 +243,7 @@ async def _download_request_with_page(self, request: Request, page: Page) -> Res
             body=body,
             request=request,
             flags=["playwright"],
+            encoding=encoding,
         )
 
     def _increment_request_stats(self, request: PlaywrightRequest) -> None:
@@ -285,3 +290,13 @@ async def _request_handler(route: Route, playwright_request: PlaywrightRequest)
             await route.continue_(**overrides)
 
         return _request_handler
+
+
+def _get_response_encoding(headers: Headers, body: str) -> Optional[str]:
+    encoding = None
+    if headers.get("content-type"):
+        content_type = to_unicode(headers["content-type"])
+        encoding = http_content_type_encoding(content_type)
+    if not encoding:
+        encoding = html_body_declared_encoding(body)
+    return encoding
diff --git a/tests/test_misc.py b/tests/test_misc.py
@@ -0,0 +1,29 @@
+import pytest
+from scrapy.http.headers import Headers
+
+from scrapy_playwright.handler import _get_response_encoding
+
+
+@pytest.mark.asyncio
+async def test_get_response_encoding():
+    assert (
+        _get_response_encoding(
+            headers=Headers({"content-type": "text/html; charset=UTF-8"}),
+            body="",
+        )
+        == "utf-8"
+    )
+    assert (
+        _get_response_encoding(
+            headers=Headers(),
+            body="""<!doctype html>
+<html lang="cn">
+<head>
+  <meta charset="gb2312">
+</head>
+</html>
+""",
+        )
+        == "gb18030"
+    )
+    assert _get_response_encoding(headers=Headers(), body="") is None
