Skip to content

Commit

Permalink
Add SoupHttpResponse class for type hinting
Browse files Browse the repository at this point in the history
  • Loading branch information
Stormheg committed Jul 27, 2023
1 parent 466b578 commit 4c956d3
Showing 1 changed file with 26 additions and 3 deletions.
29 changes: 26 additions & 3 deletions wagtail/test/utils/wagtail_tests.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,34 @@
import warnings
from contextlib import contextmanager

from bs4 import BeautifulSoup
from bs4 import BeautifulSoup, ResultSet, Tag
from django import VERSION as DJANGO_VERSION
from django.contrib.auth import get_user_model
from django.http import HttpResponse
from django.test.client import Client as DjangoTestClient
from django.test.testcases import assert_and_parse_html
from django.utils.functional import SimpleLazyObject


class SoupHttpResponse(HttpResponse):
"""Subclass for type-hinting purposes.
Use it like this for type-hinting:
response: SoupHttpResponse = self.client.get("/your-url/")
# Wonderful type hinting!
response.soup.select(...)
"""

soup: BeautifulSoup

def select(self, selector: str) -> ResultSet[Tag]:
pass

def select_one(self, selector: str) -> (Tag | None):
pass


class BS4SoupClient(DjangoTestClient):
"""Convenience wrapper around Django's Client and BeautifulSoup's SoupSieve
Expand All @@ -31,10 +51,10 @@ def _get_soup(self, response):
if "text/html" in content_type:
return BeautifulSoup(response.content, "html5lib")
raise ValueError(
"Cannot use BeatifulSoup, the response content type is not HTML"
"Cannot use BeautifulSoup, the response content type is not HTML"
)

def request(self, **request):
def request(self, **request) -> SoupHttpResponse:
response = super().request(**request)
# Make the soup available lazily. We will only incur the cost of instantiating
# the soup upon the initial access of this attribute.
Expand All @@ -49,6 +69,9 @@ def request(self, **request):
class WagtailTestUtils:
client_class = BS4SoupClient

# To make IDE type hinting recognize the appropriate type
client: BS4SoupClient

@staticmethod
def create_test_user():
"""
Expand Down

0 comments on commit 4c956d3

Please sign in to comment.