From 5ac2aa84526e70eecc34fb38cdfd7ade0f909c53 Mon Sep 17 00:00:00 2001
From: annbgn <fake@mail.com>
Date: Wed, 14 Jul 2021 00:46:45 +0300
Subject: [PATCH]  add suppport for `[of S]`? part in nth-child's arguments

---
 cssselect/parser.py     | 12 +++++++++---
 cssselect/xpath.py      |  6 ++++--
 tests/test_cssselect.py |  9 +++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/cssselect/parser.py b/cssselect/parser.py
index 7125030..0447fef 100644
--- a/cssselect/parser.py
+++ b/cssselect/parser.py
@@ -614,12 +614,13 @@ def parse_series(tokens):
     Parses the arguments for :nth-child() and friends.
 
     :raises: A list of tokens
-    :returns: :``(a, b)``
+    :returns: :``(a, b, c?)``
 
     """
     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 +631,8 @@ def parse_series(tokens):
     if 'n' not in s:
         # Just b
         return 0, int(s)
-    a, b = s.split('n', 1)
+    s, *subselector = s.split("of")
+    a, b = s.split("n", 1)
     if not a:
         a = 1
     elif a == '-' or a == '+':
@@ -641,7 +643,11 @@ def parse_series(tokens):
         b = 0
     else:
         b = int(b)
-    return a, b
+    if not subselector:
+        c = 0
+    else:
+        c = subselector[0]
+    return a, b, c
 
 
 #### Token objects
diff --git a/cssselect/xpath.py b/cssselect/xpath.py
index a8722bb..461b919 100644
--- a/cssselect/xpath.py
+++ b/cssselect/xpath.py
@@ -375,7 +375,7 @@ def xpath_indirect_adjacent_combinator(self, left, right):
     def xpath_nth_child_function(self, xpath, function, last=False,
                                  add_name_test=True):
         try:
-            a, b = parse_series(function.arguments)
+            a, b, c = parse_series(function.arguments)
         except ValueError:
             raise ExpressionError("Invalid series: '%r'" % function.arguments)
 
@@ -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 c:
+            nodetest = self.xpath(parse(c)[0].parsed_tree).element
+        elif add_name_test:
             nodetest = '*'
         else:
             nodetest  = '%s' % xpath.element
diff --git a/tests/test_cssselect.py b/tests/test_cssselect.py
index d6969f2..2027df7 100644
--- a/tests/test_cssselect.py
+++ b/tests/test_cssselect.py
@@ -406,6 +406,15 @@ 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::e[@class and contains(concat(' ', normalize-space(@class), ' '), ' important ')]) mod 2 = 0"
         assert xpath('e:nth-child(1)') == (
             "e[count(preceding-sibling::*) = 0]")