Use new WeasyPrint URL fetcher

Fix #35.

Author: liZe

flask_weasyprint/__init__.py

@@ -2,9 +2,10 @@
 
 from io import BytesIO
 from urllib.parse import urljoin, urlsplit
+from urllib.request import BaseHandler
 
 from flask import current_app, has_request_context, request, send_file
-from werkzeug.test import Client, ClientRedirectError, EnvironBuilder
+from werkzeug.test import Client, EnvironBuilder
 from werkzeug.wrappers import Response
 
 VERSION = __version__ = '1.1.0'
@@ -73,7 +74,7 @@ def dispatch(url_string):
 
 
 def make_url_fetcher(dispatcher=None, next_fetcher=True):
-    """Return an function suitable as a ``url_fetcher`` in WeasyPrint.
+    """Return a URL fetcher that handles the Flask app routes internally.
 
     You generally don’t need to call this directly.
 
@@ -92,45 +93,44 @@ def make_url_fetcher(dispatcher=None, next_fetcher=True):
     Typically ``base_url + path`` is equivalent to the passed URL.
 
     """
+    from weasyprint.urls import URLFetcher, URLFetcherResponse  # lazy loading
+
     if next_fetcher is True:
-        from weasyprint import default_url_fetcher  # lazy loading
-        next_fetcher = default_url_fetcher
+        next_fetcher = URLFetcher
 
     if dispatcher is None:
         dispatcher = make_flask_url_dispatcher()
 
-    def flask_url_fetcher(url):
-        redirect_chain = set()
-        while True:
-            result = dispatcher(url)
-            if result is None:
-                return next_fetcher(url)
-            app, base_url, path = result
-            client = Client(app, response_wrapper=Response)
-            if has_request_context() and request.cookies:
-                server_name = EnvironBuilder(
-                    path, base_url=base_url).server_name
-                for cookie_key, cookie_value in request.cookies.items():
-                    client.set_cookie(
-                        cookie_key, cookie_value, domain=server_name)
-            response = client.get(path, base_url=base_url)
-            if response.status_code == 200:
-                return {
-                    'string': response.data, 'mime_type': response.mimetype,
-                    'encoding': 'utf-8', 'redirected_url': url}
-            # The test client can follow redirects, but do it ourselves
-            # to get access to the redirected URL.
-            elif response.status_code in (301, 302, 303, 305, 307, 308):
-                redirect_chain.add(url)
-                url = urljoin(url, response.location)
-                if url in redirect_chain:
-                    raise ClientRedirectError('loop detected')
-            else:
-                raise ValueError(
-                    'Flask-WeasyPrint got HTTP status '
-                    f'{response.status} for {urljoin(base_url, path)}')
-
-    return flask_url_fetcher
+    class FlaskHandler(BaseHandler):
+        def default_open(self, req):
+            url = req.full_url
+            if result := dispatcher(url):
+                app, base_url, path = result
+                client = Client(app, response_wrapper=Response)
+                if has_request_context() and request.cookies:
+                    server_name = EnvironBuilder(path, base_url=base_url).server_name
+                    for cookie_key, cookie_value in request.cookies.items():
+                        client.set_cookie(cookie_key, cookie_value, domain=server_name)
+                response = client.get(path, base_url=base_url)
+                response = URLFetcherResponse(
+                        url, response.data, response.headers, response.status_code)
+                response.msg = ''
+                return response
+
+    class FlaskFetcher(next_fetcher or URLFetcher):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.add_handler(FlaskHandler())
+
+        def fetch(self, url, headers=None):
+            if dispatcher(url) is None:
+                if next_fetcher:
+                    return super().fetch(url, headers)
+                else:
+                    raise ValueError(f'Unknown Flask app URL: {url}')
+            return URLFetcher.fetch(self, url, headers)
+
+    return FlaskFetcher()
 
 
 def _wrapper(class_, *args, **kwargs):

pyproject.toml

@@ -13,7 +13,7 @@ readme = {file = 'README.rst', content-type = 'text/x-rst'}
 license = {file = 'LICENSE'}
 dependencies = [
   'flask >=2.3.0',
-  'weasyprint >=53.0',
+  'weasyprint >=68.0',
 ]
 classifiers = [
   'Development Status :: 5 - Production/Stable',

tests/test_flask_weasyprint.py

@@ -1,9 +1,11 @@
 """Tests for Flask-WeasyPrint."""
 
+from urllib.error import HTTPError
+
 import pytest
 from flask import Flask, json, jsonify, redirect, request
 from weasyprint import __version__ as weasyprint_version
-from werkzeug.test import ClientRedirectError
+from weasyprint.urls import URLFetcher, URLFetcherResponse
 
 from flask_weasyprint import CSS, HTML, make_url_fetcher, render_pdf
 
@@ -19,25 +21,25 @@ def test_url_fetcher():
     with app.test_request_context(base_url='http://example.org/bar/'):
         fetcher = make_url_fetcher()
 
-    result = fetcher('http://example.org/bar/')
-    assert result['string'].strip().startswith(b'<html>')
-    assert result['mime_type'] == 'text/html'
-    assert result['encoding'] == 'utf-8'
-    assert result['redirected_url'] == 'http://example.org/bar/foo/'
+    response = fetcher('http://example.org/bar/')
+    assert response.read().strip().startswith(b'<html>')
+    assert response.content_type == 'text/html'
+    assert response.charset == 'utf-8'
+    assert response.url == 'http://example.org/bar/foo/'
 
-    result = fetcher('http://example.org/bar/foo/graph?data=1&labels=A')
-    assert result['string'].strip().startswith(b'<svg xmlns=')
-    assert result['mime_type'] == 'image/svg+xml'
+    response = fetcher('http://example.org/bar/foo/graph?data=1&labels=A')
+    assert response.read().strip().startswith(b'<svg xmlns=')
+    assert response.content_type == 'image/svg+xml'
 
     # Also works with a custom dispatcher
     def custom_dispatcher(url_string):
         return app, 'http://example.org/bar/', '/foo/graph?data=1&labels=A'
     with app.test_request_context(base_url='http://example.org/bar/'):
         fetcher = make_url_fetcher(dispatcher=custom_dispatcher)
 
-    result = fetcher('test://')
-    assert result['string'].strip().startswith(b'<svg xmlns=')
-    assert result['mime_type'] == 'image/svg+xml'
+    response = fetcher('test://')
+    assert response.read().strip().startswith(b'<svg xmlns=')
+    assert response.content_type == 'image/svg+xml'
 
 
 def test_wrappers():
@@ -103,12 +105,14 @@ def add_redirect(old_url, new_url):
 
     with app.test_request_context():
         fetcher = make_url_fetcher()
-    result = fetcher('http://localhost/a')
-    assert result['string'] == b'Ok'
-    assert result['redirected_url'] == 'http://localhost/d'
-    with pytest.raises(ClientRedirectError):
+    response = fetcher('http://localhost/a')
+    assert response.read() == b'Ok'
+    assert response.url == 'http://localhost/d'
+    # TODO: keep HTTPError with WeasyPrint > 68.1, see #2686.
+    # with pytest.raises(HTTPError, match='infinite loop'):
+    with pytest.raises((HTTPError, RecursionError)):
         fetcher('http://localhost/1')
-    with pytest.raises(ValueError, match='404'):
+    with pytest.raises(HTTPError, match='404'):
         fetcher('http://localhost/nonexistent')
 
 
@@ -125,21 +129,22 @@ def catchall(subdomain='', path=None):
         app = [subdomain, request.script_root, request.path, query_string]
         return jsonify(app=app)
 
-    def dummy_fetcher(url):
-        return {'string': 'dummy ' + url}
+    class DummyFetcher(URLFetcher):
+        def fetch(self, url, headers=None):
+            return URLFetcherResponse(url, f'dummy {url}')
 
     def assert_app(url, host, script_root, path, query_string=''):
         """The URL was dispatched to the app with these parameters."""
-        assert json.loads(dispatcher(url)['string']) == {
+        assert json.loads(dispatcher(url).read()) == {
             'app': [host, script_root, path, query_string]}
 
     def assert_dummy(url):
         """The URL was not dispatched, the default fetcher was used."""
-        assert dispatcher(url)['string'] == 'dummy ' + url
+        assert dispatcher(url).read() == f'dummy {url}'.encode()
 
     # No SERVER_NAME config, default port
     with app.test_request_context(base_url='http://a.net/b/'):
-        dispatcher = make_url_fetcher(next_fetcher=dummy_fetcher)
+        dispatcher = make_url_fetcher(next_fetcher=DummyFetcher)
     assert_app('http://a.net/b', '', '/b', '/')
     assert_app('http://a.net/b/', '', '/b', '/')
     assert_app('http://a.net/b/', '', '/b', '/')
@@ -153,7 +158,7 @@ def assert_dummy(url):
 
     # No SERVER_NAME config, explicit default port
     with app.test_request_context(base_url='http://a.net:80/b/'):
-        dispatcher = make_url_fetcher(next_fetcher=dummy_fetcher)
+        dispatcher = make_url_fetcher(next_fetcher=DummyFetcher)
     assert_app('http://a.net/b', '', '/b', '/')
     assert_app('http://a.net/b/', '', '/b', '/')
     assert_app('http://a.net/b/c/d?e', '', '/b', '/c/d', 'e')
@@ -166,7 +171,7 @@ def assert_dummy(url):
 
     # Change the context’s port number
     with app.test_request_context(base_url='http://a.net:8888/b/'):
-        dispatcher = make_url_fetcher(next_fetcher=dummy_fetcher)
+        dispatcher = make_url_fetcher(next_fetcher=DummyFetcher)
     assert_app('http://a.net:8888/b', '', '/b', '/')
     assert_app('http://a.net:8888/b/', '', '/b', '/')
     assert_app('http://a.net:8888/b/cd?e', '', '/b', '/cd', 'e')
@@ -181,7 +186,7 @@ def assert_dummy(url):
     # Add a SERVER_NAME config
     app.config['SERVER_NAME'] = 'a.net'
     with app.test_request_context():
-        dispatcher = make_url_fetcher(next_fetcher=dummy_fetcher)
+        dispatcher = make_url_fetcher(next_fetcher=DummyFetcher)
     assert_app('http://a.net', '', '', '/')
     assert_app('http://a.net/', '', '', '/')
     assert_app('http://a.net/b/c/d?e', '', '', '/b/c/d', 'e')
@@ -195,7 +200,7 @@ def assert_dummy(url):
     # SERVER_NAME with a port number
     app.config['SERVER_NAME'] = 'a.net:8888'
     with app.test_request_context():
-        dispatcher = make_url_fetcher(next_fetcher=dummy_fetcher)
+        dispatcher = make_url_fetcher(next_fetcher=DummyFetcher)
     assert_app('http://a.net:8888', '', '', '/')
     assert_app('http://a.net:8888/', '', '', '/')
     assert_app('http://a.net:8888/b/c/d?e', '', '', '/b/c/d', 'e')
@@ -208,11 +213,9 @@ def assert_dummy(url):
 
 @pytest.mark.parametrize('url', [
     'http://example.net/Unïĉodé/pass !',
-    'http://example.net/Unïĉodé/pass !'.encode(),
     'http://example.net/foo%20bar/p%61ss%C2%A0!',
-    b'http://example.net/foo%20bar/p%61ss%C2%A0!',
 ])
 def test_funky_urls(url):
     with app.test_request_context(base_url='http://example.net/'):
         fetcher = make_url_fetcher()
-    assert fetcher(url)['string'] == 'pass !'.encode()
+    assert fetcher(url).read() == 'pass !'.encode()