Implement nth_child

Author: spicychickensauce

c_src/lazy_html.cpp

@@ -734,6 +734,36 @@ ExLazyHTML parent_node(ErlNifEnv *env, ExLazyHTML ex_lazy_html) {
 }
 FINE_NIF(parent_node, ERL_NIF_DIRTY_JOB_CPU_BOUND);
 
+std::vector<int64_t> nth_child(ErlNifEnv *env, ExLazyHTML ex_lazy_html) {
+  auto values = std::vector<int64_t>();
+  for (auto node : ex_lazy_html.resource->nodes) {
+    if (node->type != LXB_DOM_NODE_TYPE_ELEMENT) {
+      continue;
+    }
+
+    auto parent = node->parent;
+    if (parent == NULL) {
+      // We're at the root, nth_child is 1
+      values.push_back(1);
+    } else {
+      int64_t i = 1;
+      for (auto child = lxb_dom_node_first_child(parent); child != NULL;
+           child = lxb_dom_node_next(child)) {
+        if (child == node) {
+          break;
+        }
+        if (child->type == LXB_DOM_NODE_TYPE_ELEMENT) {
+          i++;
+        }
+      }
+      values.push_back(i);
+    }
+  }
+
+  return values;
+}
+FINE_NIF(nth_child, ERL_NIF_DIRTY_JOB_CPU_BOUND);
+
 std::string text(ErlNifEnv *env, ExLazyHTML ex_lazy_html) {
   auto document = ex_lazy_html.resource->document_ref->document;
 

lib/lazy_html.ex

@@ -387,6 +387,23 @@ defmodule LazyHTML do
     LazyHTML.NIF.parent_node(lazy_html)
   end
 
+  @doc """
+  Returns the positions of the selcted nodes among their siblings.
+  Note that the numbering is 1 based and doesn't include text nodes,
+  as it matches the `nth-child` CSS selector.
+
+  ## Examples
+
+      iex> lazy_html = LazyHTML.from_fragment(~S|<div><span>1</span><span>2</span></div>|)
+      iex> spans = LazyHTML.query(lazy_html, "span")
+      iex> LazyHTML.nth_child(spans)
+      [1, 2]
+  """
+  @spec nth_child(t()) :: list(integer())
+  def nth_child(lazy_html) do
+    LazyHTML.NIF.nth_child(lazy_html)
+  end
+
   @doc """
   Returns the text content of all nodes in `lazy_html`.
 

lib/lazy_html/nif.ex

@@ -22,6 +22,7 @@ defmodule LazyHTML.NIF do
   def query_by_id(_lazy_html, _id), do: err!()
   def child_nodes(_lazy_html), do: err!()
   def parent_node(_lazy_html), do: err!()
+  def nth_child(_lazy_html), do: err!()
   def text(_lazy_html), do: err!()
   def attribute(_lazy_html, _name), do: err!()
   def attributes(_lazy_html), do: err!()

test/lazy_html_test.exs

@@ -341,6 +341,27 @@ defmodule LazyHTMLTest do
     end
   end
 
+  describe "nth_child/1" do
+    test "nth_child gives position" do
+      lazy_html =
+        LazyHTML.from_fragment("""
+        <div>
+          Text isn't counted.
+          <span>1</span>
+          <!-- neither are comments -->
+          <span>2</span>
+        </div>
+        """)
+        assert LazyHTML.nth_child(lazy_html) == [1]
+        assert lazy_html["div"] |> LazyHTML.nth_child() == [1]
+        assert lazy_html["span"] |> LazyHTML.nth_child() == [1, 2]
+
+        # Verify numbering matches css selector
+        assert lazy_html["span:nth-child(1)"] |> LazyHTML.text == "1"
+        assert lazy_html["span:nth-child(2)"] |> LazyHTML.text == "2"
+    end
+  end
+
   describe "query_by_id/2" do
     test "raises when an empty id is given" do
       assert_raise ArgumentError, ~r/id cannot be empty/, fn ->