From b9effde875a7b9058524eb4638e34e80ad0ea7ac Mon Sep 17 00:00:00 2001 From: Akshay Sharma <42249933+akshaysharmajs@users.noreply.github.com> Date: Mon, 20 Jun 2022 20:15:33 -0400 Subject: [PATCH] Refactored _guess_content_type --- scrapy/responsetypes.py | 89 ++++++++++++++++++++++++++++--------- tests/test_responsetypes.py | 30 ++++++------- 2 files changed, 82 insertions(+), 37 deletions(-) diff --git a/scrapy/responsetypes.py b/scrapy/responsetypes.py index 998c25569..816b1f040 100644 --- a/scrapy/responsetypes.py +++ b/scrapy/responsetypes.py @@ -3,8 +3,10 @@ This module implements a class which returns the appropriate Response class based on different criteria. """ from mimetypes import MimeTypes +from operator import truediv from pkgutil import get_data from io import StringIO +from traceback import format_exception_only from urllib.parse import urlparse from warnings import warn @@ -48,6 +50,12 @@ class ResponseTypes: self.mimetypes.readfp(StringIO(mimedata)) for mimetype, cls in self.CLASSES.items(): self.classes[mimetype] = load_object(cls) + + self.prioritized_mime_type_checkers = ( + is_html_mime_type, + is_xml_mime_type, + self._is_text_mime_type, + ) def from_mimetype(self, mimetype): """Return the most appropriate Response class for the given mimetype""" @@ -122,6 +130,20 @@ class ResponseTypes: return self.from_mimetype('text/xml') else: return self.from_mimetype('text') + + def _is_text_mime_type(self, mime_type): + if ( + mime_type.startswith(b"text/") + or is_json_mime_type(mime_type) + or is_javascript_mime_type(mime_type) + or mime_type in ( + b"application/x-json", + b"application/json-amazonui-streaming", + b"application/x-javascript", + ) + ): + return True + return False def _guess_response_type(self, mime_type): if not mime_type: @@ -142,30 +164,52 @@ class ResponseTypes: ): return TextResponse return Response - - def _guess_content_type(self, body=None, headers=None, url=None, filename=None): - mimetype = None - - if headers and b'Content-Type' in headers: - mimetype = tuple(headers.getlist(b'Content-Type')) - else: - if headers and b'Content-Disposition' in headers: - filename = headers.get(b'Content-Disposition').split(b';')[-1].split(b'=')[-1].strip(b'"\'').decode() - elif url: - filename = url - - if filename: - mimetype, encoding = self.mimetypes.guess_type(filename) - if encoding: - mimetype = (f"application/{encoding}".encode(),) - elif mimetype: - mimetype = (mimetype.encode(),) - - if mimetype and self._guess_response_type(mimetype[-1]) is Response: - return None + + def _guess_type(self, filename=None): + mimetype, encoding = self.mimetypes.guess_type(filename) + if encoding: + mimetype = f"application/{encoding}".encode() + elif mimetype: + mimetype = mimetype.encode() return mimetype + def _guess_content_type(self, body=None, headers=None, url=None, filename=None): + if headers and b'Content-Type' in headers: + content_type_mime_type = headers.getlist(b'Content-Type')[-1] + else: + content_type_mime_type = None + + if headers and b'Content-Disposition' in headers: + content_disposition_mime_type = headers.get(b'Content-Disposition').split(b';')[-1].split(b'=')[-1].strip(b'"\'').decode() + content_disposition_mime_type = self._guess_type(content_disposition_mime_type) + else: + content_disposition_mime_type = None + + url_mime_type = self._guess_type(url) if url else None + filename_mime_type = self._guess_type(filename) if filename else None + + candidate_mime_types = ( + content_type_mime_type, + content_disposition_mime_type, + url_mime_type, + filename_mime_type, + ) + + for mime_type_checker in self.prioritized_mime_type_checkers: + for candidate_mime_type in candidate_mime_types: + if ( + candidate_mime_type is not None + and mime_type_checker(candidate_mime_type) + ): + return candidate_mime_type + return ( + content_type_mime_type + or content_disposition_mime_type + or url_mime_type + or filename_mime_type + ) + def _remove_nul_byte_from_text(self, text): """Return the text with removed null byte (b"\x00") if there are no other binary bytes in the text, otherwise return the text as-is. @@ -183,7 +227,8 @@ class ResponseTypes: body = body or b'' body = self._remove_nul_byte_from_text(body[:RESOURCE_HEADER_BUFFER_LENGTH]) http_origin = not url or urlparse(url).scheme in ("http", "https") - content_types = self._guess_content_type(body=body, headers=headers, url=url, filename=filename) + content_types = (self._guess_content_type(body=body, headers=headers, url=url, filename=filename),) + content_types = None if content_types == (None,) else content_types mime_type = extract_mime(body, content_types=content_types, http_origin=http_origin) return self._guess_response_type(mime_type) diff --git a/tests/test_responsetypes.py b/tests/test_responsetypes.py index 16ce36233..c309b75f0 100644 --- a/tests/test_responsetypes.py +++ b/tests/test_responsetypes.py @@ -94,6 +94,20 @@ class ResponseTypesTest(unittest.TestCase): ({'headers': Headers({'Content-Type': ['application/x-json']})}, TextResponse), ({'headers': Headers({'Content-Type': ['application/x-javascript']})}, TextResponse), ({'headers': Headers({'Content-Type': ['application/json-amazonui-streaming']})}, TextResponse), + ({'headers': Headers({'Content-Disposition': ['attachment; filename="data.xml.gz"']}), + 'url': 'http://www.example.com/page/'}, Response), + ({'headers': Headers({'Content-Type': ['application/pdf']})}, Response), + ({'url': 'http://www.example.com/page/file.html', + 'headers': Headers({'Content-Type': 'application/octet-stream'})}, HtmlResponse), + ({'url': 'http://www.example.com/item/file.xml', + 'headers': Headers({'Content-Type': 'application/octet-stream'})}, XmlResponse), + ({'url': 'http://www.example.com/item/file.html', + 'headers': Headers({'Content-Type': 'application/octet-stream'})}, HtmlResponse), + ({'url': 'http://www.example.com/item/file.xml', + 'headers': Headers({'Content-Disposition': ['attachment; filename="data.xml.gz"'], + 'Content-Type': 'application/octet-stream'})}, XmlResponse), + ({'url': 'http://www.example.com/item/file.pdf'}, Response), + ({'filename': 'file.pdf'}, Response), ] for source, cls in mappings: retcls = responsetypes.from_args(**source) @@ -103,13 +117,10 @@ class ResponseTypesTest(unittest.TestCase): """Each of the following test cases got affected after using xtractmime for MIME sniffing""" mappings = [ - ({'headers': Headers({'Content-Disposition': ['attachment; filename="data.xml.gz"']}), - 'url': 'http://www.example.com/page/'}, TextResponse), ({'body': b'\x00\x01\xff', 'url': 'http://www.example.com/item/', 'headers': Headers({'Content-Type': ['text/plain']})}, Response), ({'filename': '/tmp/temp^'}, TextResponse), ({'body': b'%PDF-1.4'}, Response), - ({'headers': Headers({'Content-Type': ['application/pdf']})}, TextResponse), ({'headers': Headers({'Content-Type': ['application/ecmascript']})}, TextResponse), ({'headers': Headers({'Content-Type': ['application/ld+json']})}, TextResponse), ({'headers': Headers({'Content-Encoding': ['zip'], 'Content-Type': ['text/html']})}, @@ -119,21 +130,10 @@ class ResponseTypesTest(unittest.TestCase): ({'body': b'Non HTML', 'headers': Headers({'Content-Encoding': ['zip'], 'Content-Type': ['text/html']})}, HtmlResponse), ({'body': b'Some plain text', 'headers': Headers({'Content-Type': 'application/octet-stream'})}, - TextResponse), - ({'url': b'http://www.example.com/page/file.html', - 'headers': Headers({'Content-Type': 'application/octet-stream'})}, TextResponse), - ({'url': 'http://www.example.com/item/file.xml', - 'headers': Headers({'Content-Type': 'application/octet-stream'})}, TextResponse), - ({'url': 'http://www.example.com/item/file.html', - 'headers': Headers({'Content-Type': 'application/octet-stream'})}, TextResponse), - ({'url': 'http://www.example.com/item/file.xml', - 'headers': Headers({'Content-Disposition': ['attachment; filename="data.xml.gz"'], - 'Content-Type': 'application/octet-stream'})}, TextResponse), + Response), ({'body': b'\x0c\x1b'}, TextResponse), ({'body': b'this is not '}, TextResponse), ({'body': b'this is not