Add Floki.traverse_and_update/2. Fix #206.

lib/floki.ex

@@ -228,6 +228,22 @@ defmodule Floki do
 
   def map(html_tree, fun), do: Finder.map(html_tree, fun)
 
+  @doc """
+  Traverses a HTML tree structure and returns a new tree structure that is the result of executing a function on all nodes. The function receives a tuple with {name, attributes, children}, and should either return a similar tuple or `nil` to delete the current node.
+
+  ## Examples
+
+      iex> html = {"div", [], ["hello"]}
+      iex> Floki.traverse_and_update(html, fn {"div", attrs, children} -> {"p", attrs, children} end)
+      {"p", [], ["hello"]}
+
+      iex> html = {"div", [], [{"span", [], ["hello"]}]}
+      iex> Floki.traverse_and_update(html, fn {"span", _attrs, _children} -> nil; tag -> tag end)
+      {"div", [], []}
+  """
+
+  defdelegate traverse_and_update(html_tree, fun), to: Floki.Traversal
+
   @doc """
   Returns the text nodes from a HTML tree.
   By default, it will perform a deep search through the HTML tree.

lib/floki/traversal.ex

@@ -0,0 +1,23 @@
+defmodule Floki.Traversal do
+  @moduledoc false
+
+  def traverse_and_update(text, _fun) when is_binary(text), do: text
+
+  def traverse_and_update([head | tail], fun) do
+    case traverse_and_update(head, fun) do
+      nil -> traverse_and_update(tail, fun)
+      mapped_head -> [mapped_head | traverse_and_update(tail, fun)]
+    end
+  end
+
+  def traverse_and_update([], _fun), do: []
+
+  def traverse_and_update({elem, attrs, children}, fun) do
+    mapped_children = traverse_and_update(children, fun)
+    fun.({elem, attrs, mapped_children})
+  end
+
+  def traverse_and_update({:comment, children}, fun) do
+    fun.({:comment, children})
+  end
+end

test/floki/traversal_test.exs

@@ -0,0 +1,58 @@
+defmodule Floki.TraversalTest do
+  use ExUnit.Case, async: true
+
+  describe "traverse_and_update/2" do
+    test "with string returns string" do
+      assert Floki.traverse_and_update("hello", fn tag -> tag end) == "hello"
+    end
+
+    test "with div and identity function returns div" do
+      html = {"div", [], ["hello"]}
+
+      assert Floki.traverse_and_update(html, fn tag -> tag end) == html
+    end
+
+    test "with div and div->p function returns p" do
+      html = {"div", [], ["hello"]}
+
+      assert Floki.traverse_and_update(html, fn {"div", attrs, children} ->
+               {"p", attrs, children}
+             end) ==
+               {"p", [], ["hello"]}
+    end
+
+    test "with style attribute and style->src function returns src attribute" do
+      html = {"div", [{"style", "display: flex"}], ["hello"]}
+
+      assert Floki.traverse_and_update(html, fn {elem, _attrs, children} ->
+               {elem, [{"src", "http://example.test"}], children}
+             end) ==
+               {"div", [{"src", "http://example.test"}], ["hello"]}
+    end
+
+    test "with text child and child replacer function returns tag with new child" do
+      html = {"div", [], ["hello"]}
+
+      assert Floki.traverse_and_update(html, fn {elem, attrs, _children} ->
+               {elem, attrs, ["world"]}
+             end) ==
+               {"div", [], ["world"]}
+    end
+
+    test "with div->span and span deleter function returns div without span" do
+      html = {"div", [], [{"span", [], ["hello"]}]}
+
+      assert Floki.traverse_and_update(html, fn
+               {"span", _attrs, _children} -> nil
+               tag -> tag
+             end) ==
+               {"div", [], []}
+    end
+
+    test "with div,p and identity function returns div,p" do
+      html = [{"div", [], ["hello"]}, {"p", [], ["world"]}]
+
+      assert Floki.traverse_and_update(html, fn tag -> tag end) == html
+    end
+  end
+end