Implement Inspect protocol for HTMLTree (#547)

* WIP Implement Inspect protocol * Fix formatting * Traverse tree only once * Change params of fun/2 recursion * Make things work * Fix formatting * Address code review improvements * Apply suggestions from code review
philss · Mar 23, 2024 · 39f4313 · 39f4313
1 parent d1fd549
commit 39f4313
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 0 deletions.
diff --git a/lib/floki/html_tree.ex b/lib/floki/html_tree.ex
@@ -328,4 +328,77 @@ defmodule Floki.HTMLTree do
       do_reduce(tree, fun.(head_node, acc), fun)
     end
   end
+
+  defimpl Inspect do
+    import Inspect.Algebra
+
+    def inspect(html_tree, opts) do
+      open = "#Floki.HTMLTree["
+      close = "]"
+      container_opts = [separator: "", break: :flex]
+
+      container_doc(
+        open,
+        nodes_with_tree(html_tree, html_tree.root_nodes_ids),
+        close,
+        opts,
+        &fun/2,
+        container_opts
+      )
+    end
+
+    defp fun({html_tree, %HTMLNode{} = html_node}, opts) do
+      {open, close, container_opts} = build_node(html_node, opts)
+
+      container_doc(
+        open,
+        nodes_with_tree(html_tree, html_node.children_nodes_ids),
+        close,
+        opts,
+        &fun/2,
+        container_opts
+      )
+    end
+
+    defp fun(%Comment{content: comment}, opts),
+      do: color(concat(["<!-- ", comment, " -->"]), :comment, opts)
+
+    defp fun(%Text{content: text}, opts), do: color(text, :string, opts)
+
+    defp nodes_with_tree(html_tree, nodes_ids) do
+      nodes_ids
+      |> Enum.reverse()
+      |> Enum.map(fn node_id ->
+        with %HTMLNode{} = html_node <- Map.get(html_tree.nodes, node_id) do
+          {html_tree, html_node}
+        end
+      end)
+    end
+
+    defp build_node(%HTMLNode{} = node, opts) do
+      tag_color = :map
+      attribute_color = :map
+
+      built_attributes =
+        for {name, value} <- node.attributes do
+          concat([
+            color(" #{name}=", attribute_color, opts),
+            color("\"#{value}\"", :string, opts)
+          ])
+        end
+        |> concat()
+
+      open =
+        concat([
+          color("<#{node.type}", tag_color, opts),
+          built_attributes,
+          color(">", tag_color, opts)
+        ])
+
+      close = color("</#{node.type}>", tag_color, opts)
+      container_opts = [separator: "", break: :strict]
+
+      {open, close, container_opts}
+    end
+  end
 end
diff --git a/test/floki/html_tree_test.exs b/test/floki/html_tree_test.exs
@@ -256,4 +256,42 @@ defmodule Floki.HTMLTreeTest do
     refute Enum.member?(html_tree, %{html_node3 | type: "marquee"})
     refute Enum.member?(html_tree, 42)
   end
+
+  test "inspect works with HTMLTree" do
+    html_tree = %HTMLTree{
+      root_nodes_ids: [1],
+      node_ids: [1],
+      nodes: %{
+        1 => %Text{content: "hello world", node_id: 1}
+      }
+    }
+
+    assert inspect(html_tree) ==
+             ~s|#Floki.HTMLTree[hello world]|
+  end
+
+  test "inspect works with HTMLTree with nested nodes" do
+    html_tree =
+      %HTMLTree{
+        root_nodes_ids: [1],
+        node_ids: [6, 5, 4, 3, 2, 1],
+        nodes: %{
+          1 => %HTMLNode{type: "html", children_nodes_ids: [6, 3, 2], node_id: 1},
+          2 => %Comment{content: "start of the stack", node_id: 2, parent_node_id: 1},
+          3 => %HTMLNode{
+            type: "a",
+            attributes: [{"class", "link"}],
+            parent_node_id: 1,
+            children_nodes_ids: [4],
+            node_id: 3
+          },
+          4 => %HTMLNode{type: "b", parent_node_id: 3, children_nodes_ids: [5], node_id: 4},
+          5 => %Text{content: "click me", parent_node_id: 4, node_id: 5},
+          6 => %HTMLNode{type: "span", parent_node_id: 1, node_id: 6}
+        }
+      }
+
+    assert inspect(html_tree) ==
+             ~s|#Floki.HTMLTree[<html><!-- start of the stack --> <a class=\"link\"><b>click me</b></a> <span></span></html>]|
+  end
 end