Add re_first() to SelectorList and iflatten() to utils.python

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]
 

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):

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 = '<ul><li id="1">1</li><li id="2">2</li></ul>'
+        response = TextResponse(url="http://example.com", body=body)
+        sel = self.sscls(response)
+
+        self.assertEqual(sel.xpath('//ul/li/text()').re_first('\d'),
+                         sel.xpath('//ul/li/text()').re('\d')[0])
+
+        self.assertEqual(sel.xpath('//ul/li[@id="1"]/text()').re_first('\d'),
+                         sel.xpath('//ul/li[@id="1"]/text()').re('\d')[0])
+
+        self.assertEqual(sel.xpath('//ul/li[2]/text()').re_first('\d'),
+                         sel.xpath('//ul/li/text()').re('\d')[1])
+
+        self.assertEqual(sel.xpath('/ul/li/text()').re_first('\w+'), None)
+        self.assertEqual(sel.xpath('/ul/li[@id="doesnt-exist"]/text()').re_first('\d'), None)
+
     def test_select_unicode_query(self):
         body = u"<p><input name='\xa9' value='1'/></p>"
         response = TextResponse(url="http://example.com", body=body, encoding='utf8')