Fix Code.Fragment.surround_context for keyword keys

Author: sabiwara

lib/elixir/lib/code/fragment.ex

@@ -640,7 +640,16 @@ defmodule Code.Fragment do
         maybe_operator(reversed_pre, post, line, opts)
 
       {:identifier, reversed_post, rest} ->
-        {rest, _} = strip_spaces(rest, 0)
+        {keyword?, rest} =
+          case rest do
+            [?: | tail] when tail == [] or hd(tail) != ?: ->
+              {true, rest}
+
+            _ ->
+              {rest, _} = strip_spaces(rest, 0)
+              {false, rest}
+          end
+
         reversed = reversed_post ++ reversed_pre
 
         case codepoint_cursor_context(reversed, opts) do
@@ -656,6 +665,9 @@ defmodule Code.Fragment do
           {{:dot, _, [_ | _]} = dot, offset} ->
             build_surround(dot, reversed, line, offset)
 
+          {{:local_or_var, acc}, offset} when keyword? ->
+            build_surround({:keyword, acc}, reversed, line, offset)
+
           {{:local_or_var, acc}, offset} when hd(rest) == ?( ->
             build_surround({:local_call, acc}, reversed, line, offset)
 

lib/elixir/test/elixir/code_fragment_test.exs

@@ -1110,6 +1110,38 @@ defmodule CodeFragmentTest do
 
       assert CF.surround_context(":", {1, 1}) == :none
     end
+
+    test "keyword keys" do
+      for i <- 2..4 do
+        assert CF.surround_context("[foo:", {1, i}) == %{
+                 context: {:keyword, ~c"foo"},
+                 begin: {1, 2},
+                 end: {1, 5}
+               }
+      end
+
+      for i <- 10..12 do
+        assert CF.surround_context("[foo: 1, bar: 2]", {1, i}) == %{
+                 context: {:keyword, ~c"bar"},
+                 begin: {1, 10},
+                 end: {1, 13}
+               }
+      end
+    end
+
+    test "keyword false positives" do
+      assert CF.surround_context("<<foo::", {1, 3}) == %{
+               context: {:local_or_var, ~c"foo"},
+               begin: {1, 3},
+               end: {1, 6}
+             }
+
+      assert CF.surround_context("[foo  :atom", {1, 2}) == %{
+               context: {:local_or_var, ~c"foo"},
+               begin: {1, 2},
+               end: {1, 5}
+             }
+    end
   end
 
   describe "container_cursor_to_quoted/2" do