diff options
| -rw-r--r-- | test/test_networking.py | 19 | ||||
| -rw-r--r-- | yt_dlp/networking/_curlcffi.py | 10 | ||||
| -rw-r--r-- | yt_dlp/networking/common.py | 6 |
3 files changed, 34 insertions, 1 deletions
diff --git a/test/test_networking.py b/test/test_networking.py index b50f70d08..d613cb568 100644 --- a/test/test_networking.py +++ b/test/test_networking.py @@ -785,6 +785,25 @@ class TestHTTPImpersonateRequestHandler(TestRequestHandlerBase): assert res.status == 200 assert std_headers['user-agent'].lower() not in res.read().decode().lower() + def test_response_extensions(self, handler): + with handler() as rh: + for target in rh.supported_targets: + request = Request( + f'http://127.0.0.1:{self.http_port}/gen_200', extensions={'impersonate': target}) + res = validate_and_send(rh, request) + assert res.extensions['impersonate'] == rh._get_request_target(request) + + def test_http_error_response_extensions(self, handler): + with handler() as rh: + for target in rh.supported_targets: + request = Request( + f'http://127.0.0.1:{self.http_port}/gen_404', extensions={'impersonate': target}) + try: + validate_and_send(rh, request) + except HTTPError as e: + res = e.response + assert res.extensions['impersonate'] == rh._get_request_target(request) + class TestRequestHandlerMisc: """Misc generic tests for request handlers, not related to request or validation testing""" diff --git a/yt_dlp/networking/_curlcffi.py b/yt_dlp/networking/_curlcffi.py index 39d1f70fb..10751a105 100644 --- a/yt_dlp/networking/_curlcffi.py +++ b/yt_dlp/networking/_curlcffi.py @@ -132,6 +132,16 @@ class CurlCFFIRH(ImpersonateRequestHandler, InstanceStoreMixin): extensions.pop('cookiejar', None) extensions.pop('timeout', None) + def send(self, request: Request) -> Response: + target = self._get_request_target(request) + try: + response = super().send(request) + except HTTPError as e: + e.response.extensions['impersonate'] = target + raise + response.extensions['impersonate'] = target + return response + def _send(self, request: Request): max_redirects_exceeded = False session: curl_cffi.requests.Session = self._get_instance( diff --git a/yt_dlp/networking/common.py b/yt_dlp/networking/common.py index 4c66ba66a..a2217034c 100644 --- a/yt_dlp/networking/common.py +++ b/yt_dlp/networking/common.py @@ -497,6 +497,7 @@ class Response(io.IOBase): @param headers: response headers. @param status: Response HTTP status code. Default is 200 OK. @param reason: HTTP status reason. Will use built-in reasons based on status code if not provided. + @param extensions: Dictionary of handler-specific response extensions. """ def __init__( @@ -505,7 +506,9 @@ class Response(io.IOBase): url: str, headers: Mapping[str, str], status: int = 200, - reason: str = None): + reason: str = None, + extensions: dict = None + ): self.fp = fp self.headers = Message() @@ -517,6 +520,7 @@ class Response(io.IOBase): self.reason = reason or HTTPStatus(status).phrase except ValueError: self.reason = None + self.extensions = extensions or {} def readable(self): return self.fp.readable() |