diff --git a/scrapy/selector/unified.py b/scrapy/selector/unified.py index 3d9435665..889c349e3 100644 --- a/scrapy/selector/unified.py +++ b/scrapy/selector/unified.py @@ -6,7 +6,7 @@ from lxml import etree from scrapy.utils.misc import extract_regex from scrapy.utils.trackref import object_ref -from scrapy.utils.python import unicode_to_str, flatten +from scrapy.utils.python import unicode_to_str, flatten, iflatten from scrapy.utils.decorator import deprecated from scrapy.http import HtmlResponse, XmlResponse from .lxmldocument import LxmlDocument @@ -175,6 +175,10 @@ class SelectorList(list): def re(self, regex): return flatten([x.re(regex) for x in self]) + def re_first(self, regex): + for el in iflatten((x.re(regex) for x in self)): + return el + def extract(self): return [x.extract() for x in self] diff --git a/scrapy/utils/python.py b/scrapy/utils/python.py index 551d337eb..b6100f899 100644 --- a/scrapy/utils/python.py +++ b/scrapy/utils/python.py @@ -27,13 +27,20 @@ def flatten(x): >>> flatten([[[1,2,3], (42,None)], [4,5], [6], 7, (8,9,10)]) [1, 2, 3, 42, None, 4, 5, 6, 7, 8, 9, 10]""" - result = [] + return list(iflatten(x)) + + +def iflatten(x): + """iflatten(sequence) -> iterator + + Similar to ``.flatten()``, but returns iterator instead""" + for el in x: if hasattr(el, "__iter__"): - result.extend(flatten(el)) + for el_ in flatten(el): + yield el_ else: - result.append(el) - return result + yield el def unique(list_, key=lambda x: x): diff --git a/tests/test_selector.py b/tests/test_selector.py index 80a9a4672..9b8613319 100644 --- a/tests/test_selector.py +++ b/tests/test_selector.py @@ -72,6 +72,24 @@ class SelectorTestCase(unittest.TestCase): self.assertEqual(sel.xpath('/ul/li[@id="doesnt-exist"]/text()').extract_first(), None) + def test_re_first(self): + """Test if re_first() returns first matched element""" + body = '