From 18e5ac49ed21e08bf996be029a9f79e53a009fff Mon Sep 17 00:00:00 2001
From: "Storm B. Heg" <storm@stormbase.digital>
Date: Thu, 27 Jul 2023 17:07:36 +0200
Subject: [PATCH] Add SoupHttpResponse class for type hinting

---
 wagtail/test/utils/wagtail_tests.py | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/wagtail/test/utils/wagtail_tests.py b/wagtail/test/utils/wagtail_tests.py
index 51499c4f4c41..fe5eed8e4f06 100644
--- a/wagtail/test/utils/wagtail_tests.py
+++ b/wagtail/test/utils/wagtail_tests.py
@@ -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
 
@@ -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.
@@ -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():
         """