Extend ANSI code parser to handle multiple arguments (#569)

* Extend ANSI code parser to handle multiple arguments

* Update changelog

CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Added support for configuring file systems using env variables ([#498](https://github.com/livebook-dev/livebook/pull/498))
 - Added a keyboard shortcut for triggering on-hover docs ([#508](https://github.com/livebook-dev/livebook/pull/508))
 - Added `--open-new` CLI flag to `livebook server` ([#529](https://github.com/livebook-dev/livebook/pull/529))
+- Nx introductory notebook ([#528](https://github.com/livebook-dev/livebook/pull/528))
 
 ### Changed
 
@@ -25,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Improved the evaluator process to not consume user-submitted messages from inbox ([#502](https://github.com/livebook-dev/livebook/pull/502))
 - Improved sections panel UI to better handle numerous sections or long section names ([#534](https://github.com/livebook-dev/livebook/pull/534) and [#537](https://github.com/livebook-dev/livebook/pull/537))
 - Fixed branching section evaluation when the parent section is empty ([#560](https://github.com/livebook-dev/livebook/pull/560)
+- Fixed ANSI support to handle multi-code escape sequences ([#569](https://github.com/livebook-dev/livebook/pull/569)
 
 ## [v0.2.3](https://github.com/livebook-dev/livebook/tree/v0.2.3) (2021-08-12)
 

lib/livebook/utils/ansi.ex

@@ -1,20 +1,6 @@
-defmodule Livebook.Utils.ANSI.Modifier do
-  @moduledoc false
-
-  defmacro defmodifier(modifier, code, terminator \\ "m") do
-    quote bind_quoted: [modifier: modifier, code: code, terminator: terminator] do
-      defp ansi_prefix_to_modifier(unquote("[#{code}#{terminator}") <> rest) do
-        {:ok, unquote(modifier), rest}
-      end
-    end
-  end
-end
-
 defmodule Livebook.Utils.ANSI do
   @moduledoc false
 
-  import Livebook.Utils.ANSI.Modifier
-
   @type modifier ::
           {:font_weight, :bold | :light}
           | {:font_style, :italic}
@@ -54,12 +40,12 @@ defmodule Livebook.Utils.ANSI do
     {tail_parts, _} =
       Enum.map_reduce(ansi_prefixed_strings, %{}, fn string, modifiers ->
         {modifiers, rest} =
-          case ansi_prefix_to_modifier(string) do
+          case ansi_prefix_to_modifiers(string) do
-            {:ok, modifier, rest} ->
+            {:ok, new_modifiers, rest} ->
-              modifiers = add_modifier(modifiers, modifier)
+              modifiers = Enum.reduce(new_modifiers, modifiers, &apply_modifier(&2, &1))
               {modifiers, rest}
 
-            {:error, _rest} ->
+            :error ->
               {modifiers, "\e" <> string}
           end
 
@@ -83,84 +69,138 @@ defmodule Livebook.Utils.ANSI do
     merge_adjacent_parts(parts, [part | acc])
   end
 
-  # Below goes a number of `ansi_prefix_to_modifier` function definitions,
+  defp ansi_prefix_to_modifiers("[1A" <> rest), do: {:ok, [:ignored], rest}
-  # that take a string like "[32msomething" (starting with ANSI code without the leading "\e")
+  defp ansi_prefix_to_modifiers("[1B" <> rest), do: {:ok, [:ignored], rest}
-  # and parse the prefix into the corresponding modifier.
+  defp ansi_prefix_to_modifiers("[1C" <> rest), do: {:ok, [:ignored], rest}
-  # The function returns either {:ok, modifier, rest} or {:error, rest}
+  defp ansi_prefix_to_modifiers("[1D" <> rest), do: {:ok, [:ignored], rest}
+  defp ansi_prefix_to_modifiers("[2J" <> rest), do: {:ok, [:ignored], rest}
+  defp ansi_prefix_to_modifiers("[2K" <> rest), do: {:ok, [:ignored], rest}
+  defp ansi_prefix_to_modifiers("[H" <> rest), do: {:ok, [:ignored], rest}
 
-  defmodifier(:reset, 0)
+  # "\e(B" is RFC1468's switch to ASCII character set and can be ignored. This
+  # can appear even when JIS character sets aren't in use
+  defp ansi_prefix_to_modifiers("(B" <> rest), do: {:ok, [:ignored], rest}
 
-  # When the code is missing (i.e., "\e[m"), it is 0 for reset.
+  defp ansi_prefix_to_modifiers("[" <> rest) do
-  defmodifier(:reset, "")
+    with [args_string, rest] <- String.split(rest, "m", parts: 2),
+         {:ok, args} <- parse_ansi_args(args_string),
+         {:ok, modifiers} <- ansi_args_to_modifiers(args, []) do
+      {:ok, modifiers, rest}
+    else
+      _ -> :error
+    end
+  end
+
+  defp ansi_prefix_to_modifiers(_string), do: :error
+
+  defp parse_ansi_args(args_string) do
+    args_string
+    |> String.split(";")
+    |> Enum.reduce_while([], fn arg, parsed ->
+      case parse_ansi_arg(arg) do
+        {:ok, n} -> {:cont, [n | parsed]}
+        :error -> {:halt, :error}
+      end
+    end)
+    |> case do
+      :error -> :error
+      parsed -> {:ok, Enum.reverse(parsed)}
+    end
+  end
+
+  defp parse_ansi_arg(""), do: {:ok, 0}
+
+  defp parse_ansi_arg(string) do
+    case Integer.parse(string) do
+      {n, ""} -> {:ok, n}
+      _ -> :error
+    end
+  end
+
+  defp ansi_args_to_modifiers([], acc), do: {:ok, Enum.reverse(acc)}
+
+  defp ansi_args_to_modifiers(args, acc) do
+    case ansi_args_to_modifier(args) do
+      {:ok, modifier, args} -> ansi_args_to_modifiers(args, [modifier | acc])
+      :error -> :error
+    end
+  end
 
   @colors [:black, :red, :green, :yellow, :blue, :magenta, :cyan, :white]
 
-  for {color, index} <- Enum.with_index(@colors) do
+  defp ansi_args_to_modifier(args) do
-    defmodifier({:foreground_color, color}, 30 + index)
+    case args do
-    defmodifier({:background_color, color}, 40 + index)
+      [0 | args] ->
-    defmodifier({:foreground_color, :"light_#{color}"}, 90 + index)
+        {:ok, :reset, args}
-    defmodifier({:background_color, :"light_#{color}"}, 100 + index)
-  end
 
-  defmodifier({:foreground_color, :reset}, 39)
+      [1 | args] ->
-  defmodifier({:background_color, :reset}, 49)
+        {:ok, {:font_weight, :bold}, args}
 
-  defmodifier({:font_weight, :bold}, 1)
+      [2 | args] ->
-  defmodifier({:font_weight, :light}, 2)
+        {:ok, {:font_weight, :light}, args}
-  defmodifier({:font_style, :italic}, 3)
-  defmodifier({:text_decoration, :underline}, 4)
-  defmodifier({:text_decoration, :line_through}, 9)
-  defmodifier({:font_weight, :reset}, 22)
-  defmodifier({:font_style, :reset}, 23)
-  defmodifier({:text_decoration, :reset}, 24)
-  defmodifier({:text_decoration, :overline}, 53)
-  defmodifier({:text_decoration, :reset}, 55)
 
-  defp ansi_prefix_to_modifier("[38;5;" <> string) do
+      [3 | args] ->
-    with {:ok, color, rest} <- bit8_prefix_to_color(string) do
+        {:ok, {:font_style, :italic}, args}
-      {:ok, {:foreground_color, color}, rest}
-    end
-  end
 
-  defp ansi_prefix_to_modifier("[48;5;" <> string) do
+      [4 | args] ->
-    with {:ok, color, rest} <- bit8_prefix_to_color(string) do
+        {:ok, {:text_decoration, :underline}, args}
-      {:ok, {:background_color, color}, rest}
-    end
-  end
 
-  # "\e(B" is RFC1468's switch to ASCII character set and can be ignored. This
+      [9 | args] ->
-  # can appear even when JIS character sets aren't in use.
+        {:ok, {:text_decoration, :line_through}, args}
-  defp ansi_prefix_to_modifier("(B" <> rest) do
-    {:ok, :ignored, rest}
-  end
 
-  defp bit8_prefix_to_color(string) do
+      [22 | args] ->
-    case Integer.parse(string) do
+        {:ok, {:font_weight, :reset}, args}
-      {n, "m" <> rest} when n in 0..255 ->
+
-        color = color_from_code(n)
+      [23 | args] ->
-        {:ok, color, rest}
+        {:ok, {:font_style, :reset}, args}
+
+      [24 | args] ->
+        {:ok, {:text_decoration, :reset}, args}
+
+      [n | args] when n in 30..37 ->
+        color = Enum.at(@colors, n - 30)
+        {:ok, {:foreground_color, color}, args}
+
+      [38, 5, bit8 | args] when bit8 in 0..255 ->
+        color = color_from_code(bit8)
+        {:ok, {:foreground_color, color}, args}
+
+      [39 | args] ->
+        {:ok, {:foreground_color, :reset}, args}
+
+      [n | args] when n in 40..47 ->
+        color = Enum.at(@colors, n - 40)
+        {:ok, {:background_color, color}, args}
+
+      [48, 5, bit8 | args] when bit8 in 0..255 ->
+        color = color_from_code(bit8)
+        {:ok, {:background_color, color}, args}
+
+      [49 | args] ->
+        {:ok, {:background_color, :reset}, args}
+
+      [53 | args] ->
+        {:ok, {:text_decoration, :overline}, args}
+
+      [55 | args] ->
+        {:ok, {:text_decoration, :reset}, args}
+
+      [n | args] when n in 90..97 ->
+        color = Enum.at(@colors, n - 90)
+        {:ok, {:foreground_color, :"light_#{color}"}, args}
+
+      [n | args] when n in 100..107 ->
+        color = Enum.at(@colors, n - 100)
+        {:ok, {:background_color, :"light_#{color}"}, args}
+
+      [n | args] when n <= 107 ->
+        {:ok, :ignored, args}
 
       _ ->
-        {:error, string}
+        :error
     end
   end
 
-  ignored_codes = [5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 25, 27, 51, 52, 54]
-
-  for code <- ignored_codes do
-    defmodifier(:ignored, code)
-  end
-
-  defmodifier(:ignored, 1, "A")
-  defmodifier(:ignored, 1, "B")
-  defmodifier(:ignored, 1, "C")
-  defmodifier(:ignored, 1, "D")
-  defmodifier(:ignored, 2, "J")
-  defmodifier(:ignored, 2, "K")
-  defmodifier(:ignored, "", "H")
-
-  defp ansi_prefix_to_modifier(string), do: {:error, string}
-
   defp color_from_code(code) when code in 0..7 do
     Enum.at(@colors, code)
   end
@@ -184,8 +224,8 @@ defmodule Livebook.Utils.ANSI do
     {:grayscale24, level}
   end
 
-  defp add_modifier(modifiers, :ignored), do: modifiers
+  defp apply_modifier(modifiers, :ignored), do: modifiers
-  defp add_modifier(_modifiers, :reset), do: %{}
+  defp apply_modifier(_modifiers, :reset), do: %{}
-  defp add_modifier(modifiers, {key, :reset}), do: Map.delete(modifiers, key)
+  defp apply_modifier(modifiers, {key, :reset}), do: Map.delete(modifiers, key)
-  defp add_modifier(modifiers, {key, value}), do: Map.put(modifiers, key, value)
+  defp apply_modifier(modifiers, {key, value}), do: Map.put(modifiers, key, value)
 end

test/livebook/utils/ansi_test.exs

@@ -57,12 +57,24 @@ defmodule Livebook.Utils.ANSITest do
       assert ANSI.parse_ansi_string("\e[H\e[1Acat") == [{[], "cat"}]
     end
 
-    test "returns the whole string if on ANSI code detected" do
+    test "returns the whole string if no ANSI code is detected" do
       assert ANSI.parse_ansi_string("\e[300mcat") == [{[], "\e[300mcat"}]
+      assert ANSI.parse_ansi_string("\ehmmcat") == [{[], "\ehmmcat"}]
     end
 
     test "ignores RFC 1468 switch to ASCII" do
       assert ANSI.parse_ansi_string("\e(Bcat") == [{[], "cat"}]
     end
+
+    test "supports multiple codes separated by semicolon" do
+      assert ANSI.parse_ansi_string("\e[0;34mcat") == [{[foreground_color: :blue], "cat"}]
+
+      assert ANSI.parse_ansi_string("\e[34;41mcat\e[0m") ==
+               [{[background_color: :red, foreground_color: :blue], "cat"}]
+
+      # 8-bit rgb color followed by background color
+      assert ANSI.parse_ansi_string("\e[38;5;67;41mcat\e[0m") ==
+               [{[background_color: :red, foreground_color: {:rgb6, 1, 2, 3}], "cat"}]
+    end
   end
 end