From bceb0e7f5dde5b6944bb51138e85fe0a167d2d2c Mon Sep 17 00:00:00 2001
From: Steffen Deusch <steffen@deusch.me>
Date: Fri, 1 Mar 2024 21:09:20 +0100
Subject: [PATCH] Add Floki.get_by_id/2 (#548)

---
 lib/floki.ex        | 26 ++++++++++++++++++++++++++
 test/floki_test.exs | 16 ++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/lib/floki.ex b/lib/floki.ex
index de1af63c..ff8fef9e 100644
--- a/lib/floki.ex
+++ b/lib/floki.ex
@@ -289,6 +289,32 @@ defmodule Floki do
     Finder.find(html_tree_as_tuple, selector)
   end
 
+  @doc """
+  Finds the first element in an HTML tree by id.
+
+  Returns `nil` if no element is found.
+
+  This is useful when there are IDs that contain special characters that
+  are invalid when passed as is as a CSS selector.
+  It is similar to the `getElementById` method in the browser.
+
+  ## Examples
+
+      iex> {:ok, html} = Floki.parse_fragment(~s[<p><span class="hint" id="id?foo_special:chars">hello</span></p>])
+      iex> Floki.get_by_id(html, "id?foo_special:chars")
+      {"span", [{"class", "hint"}, {"id", "id?foo_special:chars"}], ["hello"]}
+      iex> Floki.get_by_id(html, "does-not-exist")
+      nil
+
+  """
+  @spec get_by_id(html_tree() | html_node(), String.t()) :: html_tree
+  def get_by_id(html_tree_as_tuple, id)
+      when is_list(html_tree_as_tuple) or is_html_node(html_tree_as_tuple) do
+    html_tree_as_tuple
+    |> Finder.find(%Floki.Selector{id: id})
+    |> List.first()
+  end
+
   @doc """
   Changes the attribute values of the elements matched by `selector`
   with the function `mutation` and returns the whole element tree.
diff --git a/test/floki_test.exs b/test/floki_test.exs
index 876716c0..a4fce9db 100644
--- a/test/floki_test.exs
+++ b/test/floki_test.exs
@@ -1594,6 +1594,22 @@ defmodule FlokiTest do
              Floki.Finder.find(html_tree, [selector_struct_1, selector_struct_2])
   end
 
+  # Floki.get_by_id/2
+
+  test "get_by_id finds element with special characters" do
+    html =
+      document!(
+        html_body(~s"""
+        <div id="my-id?with_special!char:acters">Hello</div>
+        """)
+      )
+
+    assert {"div", [{"id", "my-id?with_special!char:acters"}], ["Hello"]} =
+             Floki.get_by_id(html, "my-id?with_special!char:acters")
+
+    refute Floki.get_by_id(html, "doesn't exist")
+  end
+
   # Floki.children/2
 
   test "returns the children elements of an element including the text" do