override replace() to prevent 'zyte_api_response' attribute from being lost

Author: BurnzZ

scrapy_zyte_api/responses.py

@@ -6,10 +6,20 @@
 
 
 class ZyteAPIMixin:
-    def __init__(self, *args, zyte_api_response=None, **kwargs):
+    def __init__(self, *args, zyte_api_response: Dict = None, **kwargs):
         super().__init__(*args, **kwargs)
         self._zyte_api_response = zyte_api_response
 
+    def replace(self, *args, **kwargs):
+        """Create a new response with the same attributes except for those given
+        new values.
+
+        NOTE: This doesn't support replacing the ``zyte_api_response`` attribute.
+        """
+        instance = super().replace(*args, **kwargs)
+        instance._zyte_api_response = self.zyte_api_response
+        return instance
+
     @property
     def zyte_api_response(self) -> Dict:
         """Contains the raw API response from Zyte API.

tests/test_responses.py

@@ -79,3 +79,23 @@ def test_text_from_api_response(api_response, cls):
     assert response.certificate is None
     assert response.ip_address is None
     assert response.protocol is None
+
+
+@pytest.mark.parametrize(
+    "api_response,cls",
+    [
+        (api_response_browser, ZyteAPITextResponse),
+        (api_response_body, ZyteAPIResponse),
+    ],
+)
+def test_response_replace(api_response, cls):
+    orig_response = cls.from_api_response(api_response())
+
+    # The ``zyte_api_response`` should not be replaced.
+    new_response = orig_response.replace(zyte_api_response={"overridden": "value"})
+    assert new_response.zyte_api_response == orig_response.zyte_api_response
+
+    # It should still work the same way
+    new_response = orig_response.replace(status=404)
+    assert new_response.status == 404
+    assert new_response.zyte_api_response == orig_response.zyte_api_response