iterators.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. import re
  2. import csv
  3. import logging
  4. try:
  5. from cStringIO import StringIO as BytesIO
  6. except ImportError:
  7. from io import BytesIO
  8. from io import StringIO
  9. import six
  10. from scrapy.http import TextResponse, Response
  11. from scrapy.selector import Selector
  12. from scrapy.utils.python import re_rsearch, to_unicode
  13. logger = logging.getLogger(__name__)
  14. def xmliter(obj, nodename):
  15. """Return a iterator of Selector's over all nodes of a XML document,
  16. given the name of the node to iterate. Useful for parsing XML feeds.
  17. obj can be:
  18. - a Response object
  19. - a unicode string
  20. - a string encoded as utf-8
  21. """
  22. nodename_patt = re.escape(nodename)
  23. HEADER_START_RE = re.compile(r'^(.*?)<\s*%s(?:\s|>)' % nodename_patt, re.S)
  24. HEADER_END_RE = re.compile(r'<\s*/%s\s*>' % nodename_patt, re.S)
  25. text = _body_or_str(obj)
  26. header_start = re.search(HEADER_START_RE, text)
  27. header_start = header_start.group(1).strip() if header_start else ''
  28. header_end = re_rsearch(HEADER_END_RE, text)
  29. header_end = text[header_end[1]:].strip() if header_end else ''
  30. r = re.compile(r'<%(np)s[\s>].*?</%(np)s>' % {'np': nodename_patt}, re.DOTALL)
  31. for match in r.finditer(text):
  32. nodetext = header_start + match.group() + header_end
  33. yield Selector(text=nodetext, type='xml').xpath('//' + nodename)[0]
  34. def xmliter_lxml(obj, nodename, namespace=None, prefix='x'):
  35. from lxml import etree
  36. reader = _StreamReader(obj)
  37. tag = '{%s}%s' % (namespace, nodename) if namespace else nodename
  38. iterable = etree.iterparse(reader, tag=tag, encoding=reader.encoding)
  39. selxpath = '//' + ('%s:%s' % (prefix, nodename) if namespace else nodename)
  40. for _, node in iterable:
  41. nodetext = etree.tostring(node, encoding='unicode')
  42. node.clear()
  43. xs = Selector(text=nodetext, type='xml')
  44. if namespace:
  45. xs.register_namespace(prefix, namespace)
  46. yield xs.xpath(selxpath)[0]
  47. class _StreamReader(object):
  48. def __init__(self, obj):
  49. self._ptr = 0
  50. if isinstance(obj, Response):
  51. self._text, self.encoding = obj.body, obj.encoding
  52. else:
  53. self._text, self.encoding = obj, 'utf-8'
  54. self._is_unicode = isinstance(self._text, six.text_type)
  55. def read(self, n=65535):
  56. self.read = self._read_unicode if self._is_unicode else self._read_string
  57. return self.read(n).lstrip()
  58. def _read_string(self, n=65535):
  59. s, e = self._ptr, self._ptr + n
  60. self._ptr = e
  61. return self._text[s:e]
  62. def _read_unicode(self, n=65535):
  63. s, e = self._ptr, self._ptr + n
  64. self._ptr = e
  65. return self._text[s:e].encode('utf-8')
  66. def csviter(obj, delimiter=None, headers=None, encoding=None, quotechar=None):
  67. """ Returns an iterator of dictionaries from the given csv object
  68. obj can be:
  69. - a Response object
  70. - a unicode string
  71. - a string encoded as utf-8
  72. delimiter is the character used to separate fields on the given obj.
  73. headers is an iterable that when provided offers the keys
  74. for the returned dictionaries, if not the first row is used.
  75. quotechar is the character used to enclosure fields on the given obj.
  76. """
  77. encoding = obj.encoding if isinstance(obj, TextResponse) else encoding or 'utf-8'
  78. def row_to_unicode(row_):
  79. return [to_unicode(field, encoding) for field in row_]
  80. # Python 3 csv reader input object needs to return strings
  81. if six.PY3:
  82. lines = StringIO(_body_or_str(obj, unicode=True))
  83. else:
  84. lines = BytesIO(_body_or_str(obj, unicode=False))
  85. kwargs = {}
  86. if delimiter: kwargs["delimiter"] = delimiter
  87. if quotechar: kwargs["quotechar"] = quotechar
  88. csv_r = csv.reader(lines, **kwargs)
  89. if not headers:
  90. try:
  91. row = next(csv_r)
  92. except StopIteration:
  93. return
  94. headers = row_to_unicode(row)
  95. for row in csv_r:
  96. row = row_to_unicode(row)
  97. if len(row) != len(headers):
  98. logger.warning("ignoring row %(csvlnum)d (length: %(csvrow)d, "
  99. "should be: %(csvheader)d)",
  100. {'csvlnum': csv_r.line_num, 'csvrow': len(row),
  101. 'csvheader': len(headers)})
  102. continue
  103. else:
  104. yield dict(zip(headers, row))
  105. def _body_or_str(obj, unicode=True):
  106. expected_types = (Response, six.text_type, six.binary_type)
  107. assert isinstance(obj, expected_types), \
  108. "obj must be %s, not %s" % (
  109. " or ".join(t.__name__ for t in expected_types),
  110. type(obj).__name__)
  111. if isinstance(obj, Response):
  112. if not unicode:
  113. return obj.body
  114. elif isinstance(obj, TextResponse):
  115. return obj.text
  116. else:
  117. return obj.body.decode('utf-8')
  118. elif isinstance(obj, six.text_type):
  119. return obj if unicode else obj.encode('utf-8')
  120. else:
  121. return obj.decode('utf-8') if unicode else obj