From dbfae267bf66448259e0741c87d963df5e38307f Mon Sep 17 00:00:00 2001
From: Paul Tremberth <paul.tremberth@gmail.com>
Date: Tue, 22 Nov 2016 18:20:16 +0100
Subject: [PATCH] Support *:nth-of-type through XPath extension function

---
 parsel/csstranslator.py              | 25 ++++++++++++++++++++++++-
 parsel/selector.py                   | 17 +++++++++++++++++
 tests/test_selector_csstranslator.py |  5 +++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/parsel/csstranslator.py b/parsel/csstranslator.py
index f752f2bf..8e132120 100644
--- a/parsel/csstranslator.py
+++ b/parsel/csstranslator.py
@@ -2,7 +2,7 @@
 from cssselect import HTMLTranslator as OriginalHTMLTranslator
 from cssselect.xpath import XPathExpr as OriginalXPathExpr
 from cssselect.xpath import _unicode_safe_getattr, ExpressionError
-from cssselect.parser import FunctionalPseudoElement
+from cssselect.parser import FunctionalPseudoElement, parse_series
 
 
 class XPathExpr(OriginalXPathExpr):
@@ -89,6 +89,29 @@ def xpath_text_simple_pseudo_element(self, xpath):
         """Support selecting text nodes using ::text pseudo-element"""
         return XPathExpr.from_xpath(xpath, textnode=True)
 
+    def xpath_star_nth_of_type_function(self, xpath, function, last=False):
+        try:
+            a, b = parse_series(function.arguments)
+        except ValueError:
+            raise ExpressionError("Invalid series: '%r'" % function.arguments)
+        if not last:
+            xp_format = 'star-nth-of-type({}, {})'
+        else:
+            xp_format = 'star-nth-last-of-type({}, {})'
+        return xpath.add_condition(xp_format.format(a, b))
+
+    def xpath_nth_of_type_function(self, xpath, function):
+        if xpath.element == '*':
+            return self.xpath_star_nth_of_type_function(xpath, function)
+        return self.xpath_nth_child_function(xpath, function,
+                                             add_name_test=False)
+
+    def xpath_nth_last_of_type_function(self, xpath, function):
+        if xpath.element == '*':
+            return self.xpath_star_nth_of_type_function(xpath, function,
+                                                        last=True)
+        return self.xpath_nth_child_function(xpath, function, last=True,
+                                             add_name_test=False)
 
 class GenericTranslator(TranslatorMixin, OriginalGenericTranslator):
     pass
diff --git a/parsel/selector.py b/parsel/selector.py
index ae3c633c..83c93e4d 100644
--- a/parsel/selector.py
+++ b/parsel/selector.py
@@ -11,6 +11,23 @@
 from .csstranslator import HTMLTranslator, GenericTranslator
 
 
+def css_star_nth_of_type(context, a, b, preceding=True):
+    node = context.context_node
+    num_siblings = len(list(node.itersiblings(node.tag, preceding=preceding)))
+    if a:
+        return (num_siblings // a) == (b-1)
+    else:
+        return num_siblings == (b-1)
+
+
+def css_star_nth_last_of_type(context, a, b):
+    return css_star_nth_of_type(context, a, b, preceding=False)
+
+ns = etree.FunctionNamespace(None)
+ns['star-nth-of-type'] = css_star_nth_of_type
+ns['star-nth-last-of-type'] = css_star_nth_of_type
+
+
 class SafeXMLParser(etree.XMLParser):
     def __init__(self, *args, **kwargs):
         kwargs.setdefault('resolve_entities', False)
diff --git a/tests/test_selector_csstranslator.py b/tests/test_selector_csstranslator.py
index 83d81b60..5591165d 100644
--- a/tests/test_selector_csstranslator.py
+++ b/tests/test_selector_csstranslator.py
@@ -146,6 +146,11 @@ def test_text_pseudo_element(self):
         self.assertEqual(self.x('p::text'), [u'lorem ipsum text'])
         self.assertEqual(self.x('p ::text'), [u'lorem ipsum text', u'hi', u'there', u'guy'])
 
+    def test_text_star_nth_pseudo_class(self):
+        self.assertEqual(len(self.x('div *:nth-of-type(2)')), 5)
+        self.assertEqual(self.x('*:nth-last-of-type(7)::attr(id)'),
+                                ['checkbox-disabled-checked'])
+
     def test_attribute_function(self):
         self.assertEqual(self.x('#p-b2::attr(id)'), [u'p-b2'])
         self.assertEqual(self.x('.cool-footer::attr(class)'), [u'cool-footer'])