Support doctests with iex prompts with a line number, e.g.: iex(1)> (…

…#10)

* Support iex(1)> in doctests

* Update changelog

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Fix implementation for multiline results. Multiline results are allowed, and they can be terminated with an empty new line or another doctest.
 - Support exception expressions (`** (ModuleName) message`) in results.
 - Desired line length for doctest result now accounts for its indentation.
+- Support doctests with iex prompts with a line number, e.g.: `iex(1)>`.
 
 ## 0.1.0 (2024-02-25)
 

lib/doctest_formatter/doctest_expression.ex

@@ -3,11 +3,12 @@ defmodule DoctestFormatter.DoctestExpression do
 
   alias DoctestFormatter.Indentation
 
-  defstruct([:lines, :result, :indentation])
+  defstruct([:lines, :result, :indentation, :iex_line_number])
 
   @type t :: %__MODULE__{
           lines: [String.t()],
           result: nil | [String.t()],
-          indentation: Indentation.t()
+          indentation: Indentation.t(),
+          iex_line_number: nil | pos_integer()
         }
 end

lib/doctest_formatter/formatter.ex

@@ -74,13 +74,10 @@ defmodule DoctestFormatter.Formatter do
   end
 
   defp format_lines(chunk, opts) do
-    first_line_symbol = "iex> "
-    next_line_symbol = "...> "
-
     desired_line_length = Keyword.get(opts, :line_length, default_elixir_line_length())
 
     line_length =
-      desired_line_length - elem(chunk.indentation, 1) - String.length(first_line_symbol)
+      desired_line_length - elem(chunk.indentation, 1) - String.length(get_prompt(chunk, 0))
 
     opts = Keyword.put(opts, :line_length, line_length)
 
@@ -91,8 +88,7 @@ defmodule DoctestFormatter.Formatter do
     |> String.split("\n")
     |> Enum.with_index()
     |> Enum.map(fn {line, index} ->
-      symbol = if(index == 0, do: first_line_symbol, else: next_line_symbol)
-      Indentation.indent(symbol <> line, chunk.indentation)
+      Indentation.indent(get_prompt(chunk, index) <> line, chunk.indentation)
     end)
   end
 
@@ -126,4 +122,22 @@ defmodule DoctestFormatter.Formatter do
   defp exception_result?(string) do
     string |> String.trim() |> String.starts_with?("** (")
   end
+
+  defp get_prompt(chunk, line_index) do
+    iex_line_number =
+      if chunk.iex_line_number do
+        "(#{chunk.iex_line_number})"
+      else
+        ""
+      end
+
+    prompt_text =
+      if line_index == 0 do
+        "iex"
+      else
+        "..."
+      end
+
+    "#{prompt_text}#{iex_line_number}> "
+  end
 end

lib/doctest_formatter/parser.ex

@@ -57,12 +57,13 @@ defmodule DoctestFormatter.Parser do
         %{acc | in_doctest: false, in_doctest_result: false, chunks: chunks}
 
       start = doctest_start(line) ->
-        {:start, code} = start
+        {:start, code, iex_line_number} = start
 
         chunks = [
           %DoctestExpression{
             indentation: Indentation.detect_indentation(line),
-            lines: [code]
+            lines: [code],
+            iex_line_number: iex_line_number
           }
           | acc.chunks
         ]
@@ -106,11 +107,12 @@ defmodule DoctestFormatter.Parser do
 
         %{acc | chunks: chunks}
 
-      {:start, code} ->
+      {:start, code, iex_line_number} ->
         chunks = [
           %DoctestExpression{
             indentation: Indentation.detect_indentation(line),
-            lines: [code]
+            lines: [code],
+            iex_line_number: iex_line_number
           }
           | acc.chunks
         ]
@@ -120,25 +122,48 @@ defmodule DoctestFormatter.Parser do
   end
 
   defp doctest_start(line) do
-    case Regex.run(~r/^(\s|\t)*(iex>)\s?(.*)$/, line) do
+    # example matches:
+    # iex> foo
+    # iex>foo
+    # iex()> foo
+    #   iex(43)> foo
+    case Regex.run(~r/^(\s|\t)*(?:iex(?:\((\d*)\))*>)\s?(.*)$/, line) do
       nil ->
         nil
 
-      [_, _indentation, "iex>", code | _] ->
-        {:start, code}
+      [_, _indentation, iex_line_number, code | _] ->
+        iex_line_number =
+          case iex_line_number do
+            "" ->
+              nil
+
+            _ ->
+              String.to_integer(iex_line_number)
+          end
+
+        {:start, code, iex_line_number}
     end
   end
 
   defp doctest_continuation(line) do
-    case Regex.run(~r/^(\s|\t)*((?:(?:\.\.\.)|(?:iex))>)\s?(.*)$/, line) do
+    # example matches:
+    # iex> foo
+    # iex>foo
+    # iex()> foo
+    # iex(43)> foo
+    # ...> foo
+    # ...()> foo
+    # ...(1)> foo
+    #     ...(1)> foo
+    case Regex.run(~r/^(\s|\t)*(?:(?:(?:\.\.\.)|(?:iex))(?:\(\d*\))*>)\s?(.*)$/, line) do
       nil ->
         if String.trim(line) === "" do
           nil
         else
           {:result, line}
         end
 
-      [_, _indentation, symbol, code | _] when symbol in ["iex>", "...>"] ->
+      [_, _indentation, code | _] ->
         {:continuation, code}
     end
   end

test/doctest_formatter/formatter_test.exs

@@ -254,6 +254,31 @@ defmodule DoctestFormatter.FormatterTest do
       assert output == desired_output
     end
 
+    test "keeps line number in iex>() prompt" do
+      input =
+        """
+        defmodule Foo do
+          @doc \"""
+          iex(4)>   add 4,2
+          6
+          \"""
+        end
+        """
+
+      desired_output =
+        """
+        defmodule Foo do
+          @doc \"""
+          iex(4)> add(4, 2)
+          6
+          \"""
+        end
+        """
+
+      output = format(input, [])
+      assert output == desired_output
+    end
+
     test "multiline doctest" do
       input =
         """
@@ -332,6 +357,35 @@ defmodule DoctestFormatter.FormatterTest do
       assert output == desired_output
     end
 
+    test "multiline doctest with 'iex(n)>' gets changed to '...(n)>'" do
+      input =
+        """
+        defmodule Foo do
+          @doc \"""
+          iex(3)>  "Fizz"
+          iex()>   |> concat( "Buzz" )
+          iex()> |>     concat("Barr")
+                   "FizzBuzzBarr"
+          \"""
+        end
+        """
+
+      desired_output =
+        """
+        defmodule Foo do
+          @doc \"""
+          iex(3)> "Fizz"
+          ...(3)> |> concat("Buzz")
+          ...(3)> |> concat("Barr")
+          "FizzBuzzBarr"
+          \"""
+        end
+        """
+
+      output = format(input, [])
+      assert output == desired_output
+    end
+
     test "doctest can get split into more lines than originally" do
       input =
         """

test/doctest_formatter/parser_test.exs

@@ -339,5 +339,72 @@ defmodule DoctestFormatter.ParserTest do
                %OtherContent{lines: ["foo", ""]}
              ]
     end
+
+    test "with iex(n)>" do
+      assert parse("    iex(3)> 1 + 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 + 2"],
+                 result: ["3"],
+                 iex_line_number: 3,
+                 indentation: {:spaces, 4}
+               }
+             ]
+
+      assert parse("iex(14)> 1 +\niex(2)> 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 +", "2"],
+                 result: ["3"],
+                 iex_line_number: 14,
+                 indentation: {:spaces, 0}
+               }
+             ]
+
+      assert parse("iex(6)> 1 +\n...(6)> 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 +", "2"],
+                 result: ["3"],
+                 iex_line_number: 6,
+                 indentation: {:spaces, 0}
+               }
+             ]
+
+      assert parse("  iex(6)> 1 +\n  ...(7)> 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 +", "2"],
+                 result: ["3"],
+                 iex_line_number: 6,
+                 indentation: {:spaces, 2}
+               }
+             ]
+    end
+
+    test "iex()> counts as no line number" do
+      assert parse("iex()> 1 + 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 + 2"],
+                 result: ["3"],
+                 iex_line_number: nil,
+                 indentation: {:spaces, 0}
+               }
+             ]
+
+      assert parse("  iex()> 1 +\n  iex()> 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 +", "2"],
+                 result: ["3"],
+                 iex_line_number: nil,
+                 indentation: {:spaces, 2}
+               }
+             ]
+
+      assert parse("iex()> 1 +\n...()> 2\n3") == [
+               %DoctestExpression{
+                 lines: ["1 +", "2"],
+                 result: ["3"],
+                 iex_line_number: nil,
+                 indentation: {:spaces, 0}
+               }
+             ]
+    end
   end
 end