add suppport for [of S]? part in nth-child's arguments

Author: annbgn

cssselect/parser.py

@@ -182,12 +182,27 @@ class Function(object):
     """
     Represents selector:name(expr)
     """
-    def __init__(self, selector, name, arguments):
+
+    def __init__(self, selector, name, arguments, of_type=None):
         self.selector = selector
         self.name = ascii_lower(name)
         self.arguments = arguments
 
+        # for css4 :nth-child(An+B of Subselector)
+        try:
+            self.of_type = of_type[0]
+        except (IndexError, TypeError):
+            self.of_type = None
+
     def __repr__(self):
+        if self.of_type:
+            return "%s[%r:%s(%r of %s)]" % (
+                self.__class__.__name__,
+                self.selector,
+                self.name,
+                [token.value for token in self.arguments],
+                self.of_type.__repr__(),
+            )
         return '%s[%r:%s(%r)]' % (
             self.__class__.__name__, self.selector, self.name,
             [token.value for token in self.arguments])
@@ -539,7 +554,8 @@ def parse_simple_selector(stream, inside_negation=False):
                     raise SelectorSyntaxError("Expected ')', got %s" % (next,))
                 result = Negation(result, argument)
             else:
-                result = Function(result, ident, parse_arguments(stream))
+                arguments, of_type = parse_arguments(stream)
+                result = Function(result, ident, arguments, of_type)
         else:
             raise SelectorSyntaxError(
                 "Expected selector, got %s" % (peek,))
@@ -554,16 +570,33 @@ def parse_arguments(stream):
     while 1:
         stream.skip_whitespace()
         next = stream.next()
-        if next.type in ('IDENT', 'STRING', 'NUMBER') or next in [
-                ('DELIM', '+'), ('DELIM', '-')]:
+        if next == ("IDENT", "of"):
+            stream.skip_whitespace()
+            of_type = parse_of_type(stream)
+            return arguments, of_type
+        elif next.type in ("IDENT", "STRING", "NUMBER") or next in [
+            ("DELIM", "+"),
+            ("DELIM", "-"),
+        ]:
             arguments.append(next)
         elif next == ('DELIM', ')'):
-            return arguments
+            return arguments, None
         else:
             raise SelectorSyntaxError(
                 "Expected an argument, got %s" % (next,))
 
 
+def parse_of_type(stream):
+    subselector = ""
+    while 1:
+        next = stream.next()
+        if next == ("DELIM", ")"):
+            break
+        subselector += next.value
+    result = parse(subselector)
+    return result
+
+
 def parse_attrib(selector, stream):
     stream.skip_whitespace()
     attrib = stream.next_ident_or_star()
@@ -620,6 +653,7 @@ def parse_series(tokens):
     for token in tokens:
         if token.type == 'STRING':
             raise ValueError('String tokens not allowed in series.')
+
     s = ''.join(token.value for token in tokens).strip()
     if s == 'odd':
         return 2, 1
@@ -630,7 +664,7 @@ def parse_series(tokens):
     if 'n' not in s:
         # Just b
         return 0, int(s)
-    a, b = s.split('n', 1)
+    a, b = s.split("n", 1)
     if not a:
         a = 1
     elif a == '-' or a == '+':
@@ -641,6 +675,7 @@ def parse_series(tokens):
         b = 0
     else:
         b = int(b)
+
     return a, b
 
 

cssselect/xpath.py

@@ -439,7 +439,9 @@ def xpath_nth_child_function(self, xpath, function, last=False,
         # `add_name_test` boolean is inverted and somewhat counter-intuitive:
         #
         # nth_of_type() calls nth_child(add_name_test=False)
-        if add_name_test:
+        if function.of_type:
+            nodetest = self.xpath(function.of_type.parsed_tree)
+        elif add_name_test:
             nodetest = '*'
         else:
             nodetest  = '%s' % xpath.element

tests/test_cssselect.py

@@ -406,6 +406,19 @@ def xpath(css):
                "@hreflang = 'en' or starts-with(@hreflang, 'en-'))]")
 
         # --- nth-* and nth-last-* -------------------------------------
+        assert (
+            xpath("e:nth-child(2n+1 of S)")
+            == "e[count(preceding-sibling::S) mod 2 = 0]"
+        )
+        assert (
+            xpath("e:nth-of-type(2n+1 of S)")
+            == "e[count(preceding-sibling::S) mod 2 = 0]"
+        )
+        assert (
+            xpath('e:nth-child(2n+1 of li.important)')
+            == "e[count(preceding-sibling::li[@class and contains(concat(' ', normalize-space(@class), ' '), ' important ')]) mod 2 = 0]"
+        )
+
         assert xpath('e:nth-child(1)') == (
             "e[count(preceding-sibling::*) = 0]")