Add more DocAST helpers

Author: josevalim

lib/ex_doc/doc_ast.ex

@@ -122,6 +122,43 @@ defmodule ExDoc.DocAST do
         do: text
   end
 
+  @doc """
+  Extracts the headers which have anchors (aka ids) in them.
+  """
+  def extract_headers_with_ids(ast, headers) do
+    ast
+    |> reduce_tags([], fn {tag, attrs, inner, _}, acc ->
+      with true <- tag in headers,
+           id = Keyword.get(attrs, :id, ""),
+           text = ExDoc.DocAST.text(inner),
+           true <- id != "" and text != "" do
+        [{tag, text, id} | acc]
+      else
+        _ -> acc
+      end
+    end)
+    |> Enum.reverse()
+  end
+
+  @doc """
+  Adds an id attribute to the given headers.
+  """
+  def add_ids_to_headers(doc_ast, headers) do
+    doc_ast
+    |> map_reduce_tags(%{}, fn {tag, attrs, inner, meta} = ast, seen ->
+      if tag in headers and not Keyword.has_key?(attrs, :id) do
+        possible_id = inner |> text() |> ExDoc.Utils.text_to_id()
+        id_count = Map.get(seen, possible_id, 0)
+        actual_id = if id_count >= 1, do: "#{possible_id}-#{id_count}", else: possible_id
+        seen = Map.put(seen, possible_id, id_count + 1)
+        {{tag, [id: actual_id] ++ attrs, inner, meta}, seen}
+      else
+        {ast, seen}
+      end
+    end)
+    |> elem(0)
+  end
+
   @doc """
   Compute a synopsis from a document by looking at its first paragraph.
   """
@@ -144,14 +181,11 @@ defmodule ExDoc.DocAST do
   @doc """
   Remove ids from elements.
   """
-  def remove_ids({tag, attrs, inner, meta}),
-    do: {tag, Keyword.delete(attrs, :href), remove_ids(inner), meta}
-
-  def remove_ids(list) when is_list(list),
-    do: Enum.map(list, &remove_ids/1)
-
-  def remove_ids(other),
-    do: other
+  def remove_ids(ast) do
+    map_tags(ast, fn {tag, attrs, inner, meta} ->
+      {tag, Keyword.delete(attrs, :href), inner, meta}
+    end)
+  end
 
   @doc """
   Returns text content from the given AST.
@@ -199,57 +233,52 @@ defmodule ExDoc.DocAST do
   defp pivot([head | tail], acc, headers), do: pivot(tail, [head | acc], headers)
   defp pivot([], acc, _headers), do: Enum.reverse(acc)
 
-  def highlight(html, language, opts \\ []) do
-    do_highlight(html, language.highlight_info(), opts)
-  end
-
-  defp do_highlight(
-         {:pre, pre_attrs, [{:code, code_attrs, [code], code_meta}], pre_meta} = ast,
-         highlight_info,
-         opts
-       )
-       when is_binary(code) do
-    {lang, code_attrs} = Keyword.pop(code_attrs, :class, "")
-
-    case pick_language_and_lexer(lang, highlight_info, code) do
-      {_lang, nil, _lexer_opts} ->
-        ast
+  @doc """
+  Highlights the code blocks in the AST.
+  """
+  def highlight(ast, language, opts \\ []) do
+    highlight_info = language.highlight_info()
 
-      {lang, lexer, lexer_opts} ->
-        try do
-          Makeup.highlight_inner_html(code,
-            lexer: lexer,
-            lexer_options: lexer_opts,
-            formatter_options: opts
-          )
-        rescue
-          exception ->
-            ExDoc.Utils.warn(
-              [
-                "crashed while highlighting #{lang} snippet:\n\n",
-                ExDoc.DocAST.to_string(ast),
-                "\n\n",
-                Exception.format_banner(:error, exception, __STACKTRACE__)
-              ],
-              __STACKTRACE__
-            )
+    map_tags(ast, fn
+      {:pre, pre_attrs, [{:code, code_attrs, [code], code_meta}], pre_meta} = ast
+      when is_binary(code) ->
+        {lang, code_attrs} = Keyword.pop(code_attrs, :class, "")
 
+        case pick_language_and_lexer(lang, highlight_info, code) do
+          {_lang, nil, _lexer_opts} ->
             ast
-        else
-          highlighted ->
-            code_attrs = [class: "makeup #{lang}", translate: "no"] ++ code_attrs
-            code_meta = Map.put(code_meta, :verbatim, true)
-            {:pre, pre_attrs, [{:code, code_attrs, [highlighted], code_meta}], pre_meta}
-        end
-    end
-  end
 
-  defp do_highlight(list, highlight_info, opts) when is_list(list) do
-    Enum.map(list, &do_highlight(&1, highlight_info, opts))
-  end
+          {lang, lexer, lexer_opts} ->
+            try do
+              Makeup.highlight_inner_html(code,
+                lexer: lexer,
+                lexer_options: lexer_opts,
+                formatter_options: opts
+              )
+            rescue
+              exception ->
+                ExDoc.Utils.warn(
+                  [
+                    "crashed while highlighting #{lang} snippet:\n\n",
+                    ExDoc.DocAST.to_string(ast),
+                    "\n\n",
+                    Exception.format_banner(:error, exception, __STACKTRACE__)
+                  ],
+                  __STACKTRACE__
+                )
+
+                ast
+            else
+              highlighted ->
+                code_attrs = [class: "makeup #{lang}", translate: "no"] ++ code_attrs
+                code_meta = Map.put(code_meta, :verbatim, true)
+                {:pre, pre_attrs, [{:code, code_attrs, [highlighted], code_meta}], pre_meta}
+            end
+        end
 
-  defp do_highlight(other, _highlight_info, _opts) do
-    other
+      ast ->
+        ast
+    end)
   end
 
   defp pick_language_and_lexer("", _highlight_info, "$ " <> _) do
@@ -270,4 +299,44 @@ defmodule ExDoc.DocAST do
       :error -> {lang, nil, []}
     end
   end
+
+  ## Traversal helpers
+
+  @doc """
+  Maps the tags in the AST, first mapping children tags, then the tag itself.
+  """
+  def map_tags({tag, attrs, inner, meta}, fun),
+    do: fun.({tag, attrs, Enum.map(inner, &map_tags(&1, fun)), meta})
+
+  def map_tags(list, fun) when is_list(list),
+    do: Enum.map(list, &map_tags(&1, fun))
+
+  def map_tags(other, _fun),
+    do: other
+
+  @doc """
+  Reduces the tags in the AST, first reducing children tags, then the tag itself.
+  """
+  def reduce_tags({tag, attrs, inner, meta}, acc, fun),
+    do: fun.({tag, attrs, inner, meta}, Enum.reduce(inner, acc, &reduce_tags(&1, &2, fun)))
+
+  def reduce_tags(list, acc, fun) when is_list(list),
+    do: Enum.reduce(list, acc, &reduce_tags(&1, &2, fun))
+
+  def reduce_tags(_other, acc, _fun),
+    do: acc
+
+  @doc """
+  Map-reduces the tags in the AST, first mapping children tags, then the tag itself.
+  """
+  def map_reduce_tags({tag, attrs, inner, meta}, acc, fun) do
+    {inner, acc} = Enum.map_reduce(inner, acc, &map_reduce_tags(&1, &2, fun))
+    fun.({tag, attrs, inner, meta}, acc)
+  end
+
+  def map_reduce_tags(list, acc, fun) when is_list(list),
+    do: Enum.map_reduce(list, acc, &map_reduce_tags(&1, &2, fun))
+
+  def map_reduce_tags(other, acc, _fun),
+    do: {other, acc}
 end

lib/ex_doc/formatter/html.ex

@@ -414,7 +414,7 @@ defmodule ExDoc.Formatter.HTML do
       case extension_name(input) do
         extension when extension in ["", ".txt"] ->
           source = File.read!(input)
-          ast = [{:pre, [], "\n" <> source, %{}}]
+          ast = [{:pre, [], ["\n" <> source], %{}}]
           {source, ast}
 
         extension when extension in [".md", ".livemd", ".cheatmd"] ->

test/ex_doc/doc_ast_test.exs

@@ -126,8 +126,8 @@ defmodule ExDoc.DocASTTest do
 
     defp synopsis(markdown) do
       markdown
-      |> ExDoc.DocAST.parse!("text/markdown")
-      |> ExDoc.DocAST.synopsis()
+      |> DocAST.parse!("text/markdown")
+      |> DocAST.synopsis()
     end
   end
 
@@ -151,8 +151,36 @@ defmodule ExDoc.DocASTTest do
 
     defp extract_headers(markdown) do
       markdown
-      |> ExDoc.DocAST.parse!("text/markdown")
-      |> ExDoc.DocAST.extract_headers([:h2])
+      |> DocAST.parse!("text/markdown")
+      |> DocAST.extract_headers([:h2])
+    end
+  end
+
+  describe "headers" do
+    test "adds and extracts anchored headers" do
+      assert """
+             # h1
+
+             ## h2
+
+             ### h3 repeat
+
+             ## h2 > h3
+
+             ### h3 repeat
+
+             > ## inside `blockquote`
+             """
+             |> DocAST.parse!("text/markdown")
+             |> DocAST.add_ids_to_headers([:h2, :h3])
+             |> DocAST.extract_headers_with_ids([:h2, :h3]) ==
+               [
+                 {:h2, "h2", "h2"},
+                 {:h3, "h3 repeat", "h3-repeat"},
+                 {:h2, "h2 > h3", "h2-h3"},
+                 {:h3, "h3 repeat", "h3-repeat-1"},
+                 {:h2, "inside blockquote", "inside-blockquote"}
+               ]
     end
   end
 
@@ -196,35 +224,55 @@ defmodule ExDoc.DocASTTest do
              ```
              """) =~
                ~r{<pre><code class=\"makeup shell\" translate="no"><span class="gp unselectable">\$.*}
+
+      # Nested in another element
+      assert highlight("""
+             > ```elixir
+             > hello
+             > ```
+             """) =~
+               ~r{<blockquote><pre><code class=\"makeup elixir\" translate="no">.*}
     end
 
     defp highlight(markdown) do
       markdown
-      |> ExDoc.DocAST.parse!("text/markdown")
-      |> ExDoc.DocAST.highlight(ExDoc.Language.Elixir)
-      |> ExDoc.DocAST.to_string()
+      |> DocAST.parse!("text/markdown")
+      |> DocAST.highlight(ExDoc.Language.Elixir)
+      |> DocAST.to_string()
     end
   end
 
   describe "sectionize" do
     test "sectioninize" do
-      list = [
-        {:h1, [], ["H1"], %{}},
-        {:h2, [class: "example"], ["H2-1"], %{}},
-        {:p, [], ["p1"], %{}},
-        {:h3, [], ["H3-1"], %{}},
-        {:p, [], ["p2"], %{}},
-        {:h3, [], ["H3-2"], %{}},
-        {:p, [], ["p3"], %{}},
-        {:h3, [], ["H3-3"], %{}},
-        {:p, [], ["p4"], %{}},
-        {:h2, [], ["H2-2"], %{}},
-        {:p, [], ["p5"], %{}},
-        {:h3, [class: "last"], ["H3-1"], %{}},
-        {:p, [], ["p6"], %{}}
-      ]
-
-      assert DocAST.sectionize(list, [:h2, :h3]) ==
+      assert """
+             # H1
+
+             ## H2-1 {:class="example"}
+
+             p1
+
+             ### H3-1
+
+             p2
+
+             ### H3-2
+
+             p3
+
+             ### H3-3
+
+             p4
+
+             ## H2-2
+
+             p5
+
+             ### H3-1 {:class="last"}
+
+             p6
+             """
+             |> DocAST.parse!("text/markdown")
+             |> DocAST.sectionize([:h2, :h3]) ==
                [
                  {:h1, [], ["H1"], %{}},
                  {:section, [class: "h2 example"],