From d1f87e4f088c0757dd833d2a0841be0b830deb57 Mon Sep 17 00:00:00 2001 From: Andrey Rakhmatullin Date: Wed, 2 Aug 2023 23:11:15 +0400 Subject: [PATCH] More typing for scrapy/utils/iterators.py. --- scrapy/http/response/text.py | 6 +- scrapy/selector/unified.py | 18 ++++-- scrapy/utils/iterators.py | 106 ++++++++++++++++++++++++++-------- tests/test_utils_iterators.py | 7 ++- 4 files changed, 103 insertions(+), 34 deletions(-) diff --git a/scrapy/http/response/text.py b/scrapy/http/response/text.py index 5289f014a..f228e11c1 100644 --- a/scrapy/http/response/text.py +++ b/scrapy/http/response/text.py @@ -7,7 +7,7 @@ See documentation in docs/topics/request-response.rst import json from contextlib import suppress -from typing import Generator, Tuple +from typing import Generator, Optional, Tuple from urllib.parse import urljoin import parsel @@ -37,7 +37,7 @@ class TextResponse(Response): def __init__(self, *args, **kwargs): self._encoding = kwargs.pop("encoding", None) self._cached_benc = None - self._cached_ubody = None + self._cached_ubody: Optional[str] = None self._cached_selector = None super().__init__(*args, **kwargs) @@ -82,7 +82,7 @@ class TextResponse(Response): return self._cached_decoded_json @property - def text(self): + def text(self) -> str: """Body as unicode""" # access self.encoding before _cached_ubody to make sure # _body_inferred_encoding is called diff --git a/scrapy/selector/unified.py b/scrapy/selector/unified.py index caff79e9c..863fb6032 100644 --- a/scrapy/selector/unified.py +++ b/scrapy/selector/unified.py @@ -1,10 +1,11 @@ """ XPath selectors based on lxml """ +from typing import Any, Optional, Type, Union from parsel import Selector as _ParselSelector -from scrapy.http import HtmlResponse, XmlResponse +from scrapy.http import HtmlResponse, TextResponse, XmlResponse from scrapy.utils.python import to_bytes from scrapy.utils.trackref import object_ref @@ -13,14 +14,14 @@ __all__ = ["Selector", "SelectorList"] _NOT_SET = object() -def _st(response, st): +def _st(response: Optional[TextResponse], st: Optional[str]) -> str: if st is None: return "xml" if isinstance(response, XmlResponse) else "html" return st -def _response_from_text(text, st): - rt = XmlResponse if st == "xml" else HtmlResponse +def _response_from_text(text: Union[str, bytes], st: Optional[str]) -> TextResponse: + rt: Type[TextResponse] = XmlResponse if st == "xml" else HtmlResponse return rt(url="about:blank", encoding="utf-8", body=to_bytes(text, "utf-8")) @@ -65,7 +66,14 @@ class Selector(_ParselSelector, object_ref): __slots__ = ["response"] selectorlist_cls = SelectorList - def __init__(self, response=None, text=None, type=None, root=_NOT_SET, **kwargs): + def __init__( + self, + response: Optional[TextResponse] = None, + text: Optional[str] = None, + type: Optional[str] = None, + root: Optional[Any] = _NOT_SET, + **kwargs: Any, + ): if response is not None and text is not None: raise ValueError( f"{self.__class__.__name__}.__init__() received " diff --git a/scrapy/utils/iterators.py b/scrapy/utils/iterators.py index 170055d5e..58850b843 100644 --- a/scrapy/utils/iterators.py +++ b/scrapy/utils/iterators.py @@ -1,16 +1,37 @@ +from __future__ import annotations + import csv import logging import re from io import StringIO +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Generator, + Iterable, + List, + Literal, + Optional, + Union, + cast, + overload, +) from scrapy.http import Response, TextResponse from scrapy.selector import Selector from scrapy.utils.python import re_rsearch, to_unicode +if TYPE_CHECKING: + from lxml._types import SupportsReadClose + logger = logging.getLogger(__name__) -def xmliter(obj, nodename): +def xmliter( + obj: Union[Response, str, bytes], nodename: str +) -> Generator[Selector, Any, None]: """Return a iterator of Selector's over all nodes of a XML document, given the name of the node to iterate. Useful for parsing XML feeds. @@ -27,20 +48,22 @@ def xmliter(obj, nodename): NAMESPACE_RE = re.compile(r"((xmlns[:A-Za-z]*)=[^>\s]+)", re.S) text = _body_or_str(obj) - document_header = re.search(DOCUMENT_HEADER_RE, text) - document_header = document_header.group().strip() if document_header else "" + document_header_match = re.search(DOCUMENT_HEADER_RE, text) + document_header = ( + document_header_match.group().strip() if document_header_match else "" + ) header_end_idx = re_rsearch(HEADER_END_RE, text) header_end = text[header_end_idx[1] :].strip() if header_end_idx else "" - namespaces = {} + namespaces: Dict[str, str] = {} if header_end: for tagname in reversed(re.findall(END_TAG_RE, header_end)): + assert header_end_idx tag = re.search( rf"<\s*{tagname}.*?xmlns[:=][^>]*>", text[: header_end_idx[1]], re.S ) if tag: - namespaces.update( - reversed(x) for x in re.findall(NAMESPACE_RE, tag.group()) - ) + for x in re.findall(NAMESPACE_RE, tag.group()): + namespaces[x[1]] = x[0] r = re.compile(rf"<{nodename_patt}[\s>].*?", re.DOTALL) for match in r.finditer(text): @@ -54,12 +77,19 @@ def xmliter(obj, nodename): yield Selector(text=nodetext, type="xml") -def xmliter_lxml(obj, nodename, namespace=None, prefix="x"): +def xmliter_lxml( + obj: Union[TextResponse, str, bytes], + nodename: str, + namespace: Optional[str] = None, + prefix: str = "x", +) -> Generator[Selector, Any, None]: from lxml import etree reader = _StreamReader(obj) tag = f"{{{namespace}}}{nodename}" if namespace else nodename - iterable = etree.iterparse(reader, tag=tag, encoding=reader.encoding) + iterable = etree.iterparse( + cast(SupportsReadClose[bytes], reader), tag=tag, encoding=reader.encoding + ) selxpath = "//" + (f"{prefix}:{nodename}" if namespace else nodename) for _, node in iterable: nodetext = etree.tostring(node, encoding="unicode") @@ -71,30 +101,39 @@ def xmliter_lxml(obj, nodename, namespace=None, prefix="x"): class _StreamReader: - def __init__(self, obj): - self._ptr = 0 - if isinstance(obj, Response): + def __init__(self, obj: Union[TextResponse, str, bytes]): + self._ptr: int = 0 + self._text: Union[str, bytes] + if isinstance(obj, TextResponse): self._text, self.encoding = obj.body, obj.encoding else: self._text, self.encoding = obj, "utf-8" - self._is_unicode = isinstance(self._text, str) + self._is_unicode: bool = isinstance(self._text, str) - def read(self, n=65535): - self.read = self._read_unicode if self._is_unicode else self._read_string + def read(self, n: int = 65535) -> bytes: + self.read: Callable[[int], bytes] = ( # type: ignore[method-assign] + self._read_unicode if self._is_unicode else self._read_string + ) return self.read(n).lstrip() - def _read_string(self, n=65535): + def _read_string(self, n: int = 65535) -> bytes: s, e = self._ptr, self._ptr + n self._ptr = e - return self._text[s:e] + return cast(bytes, self._text)[s:e] - def _read_unicode(self, n=65535): + def _read_unicode(self, n: int = 65535) -> bytes: s, e = self._ptr, self._ptr + n self._ptr = e - return self._text[s:e].encode("utf-8") + return cast(str, self._text)[s:e].encode("utf-8") -def csviter(obj, delimiter=None, headers=None, encoding=None, quotechar=None): +def csviter( + obj: Union[Response, str, bytes], + delimiter: Optional[str] = None, + headers: Optional[List[str]] = None, + encoding: Optional[str] = None, + quotechar: Optional[str] = None, +) -> Generator[Dict[str, str], Any, None]: """Returns an iterator of dictionaries from the given csv object obj can be: @@ -112,12 +151,12 @@ def csviter(obj, delimiter=None, headers=None, encoding=None, quotechar=None): encoding = obj.encoding if isinstance(obj, TextResponse) else encoding or "utf-8" - def row_to_unicode(row_): + def row_to_unicode(row_: Iterable) -> List[str]: return [to_unicode(field, encoding) for field in row_] lines = StringIO(_body_or_str(obj, unicode=True)) - kwargs = {} + kwargs: Dict[str, Any] = {} if delimiter: kwargs["delimiter"] = delimiter if quotechar: @@ -147,7 +186,24 @@ def csviter(obj, delimiter=None, headers=None, encoding=None, quotechar=None): yield dict(zip(headers, row)) -def _body_or_str(obj, unicode=True): +@overload +def _body_or_str(obj: Union[Response, str, bytes]) -> str: + ... + + +@overload +def _body_or_str(obj: Union[Response, str, bytes], unicode: Literal[True]) -> str: + ... + + +@overload +def _body_or_str(obj: Union[Response, str, bytes], unicode: Literal[False]) -> bytes: + ... + + +def _body_or_str( + obj: Union[Response, str, bytes], unicode: bool = True +) -> Union[str, bytes]: expected_types = (Response, str, bytes) if not isinstance(obj, expected_types): expected_types_str = " or ".join(t.__name__ for t in expected_types) @@ -156,10 +212,10 @@ def _body_or_str(obj, unicode=True): ) if isinstance(obj, Response): if not unicode: - return obj.body + return cast(bytes, obj.body) if isinstance(obj, TextResponse): return obj.text - return obj.body.decode("utf-8") + return cast(bytes, obj.body).decode("utf-8") if isinstance(obj, str): return obj if unicode else obj.encode("utf-8") return obj.decode("utf-8") if unicode else obj diff --git a/tests/test_utils_iterators.py b/tests/test_utils_iterators.py index 3598fa0bb..5dfd7e7ac 100644 --- a/tests/test_utils_iterators.py +++ b/tests/test_utils_iterators.py @@ -1,13 +1,18 @@ +from typing import Callable, Iterable, Union + from pytest import mark from twisted.trial import unittest +from scrapy import Selector from scrapy.http import Response, TextResponse, XmlResponse from scrapy.utils.iterators import _body_or_str, csviter, xmliter, xmliter_lxml from tests import get_testdata class XmliterTestCase(unittest.TestCase): - xmliter = staticmethod(xmliter) + xmliter: Callable[ + [Union[TextResponse, str, bytes], str], Iterable[Selector] + ] = staticmethod(xmliter) def test_xmliter(self): body = b"""