More typing for scrapy/utils/iterators.py.

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

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 "

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>].*?</{nodename_patt}>", 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

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"""