Type check Map.fetch!/2 and :maps.take/2

Author: josevalim

lib/elixir/lib/module/types/apply.ex

@@ -245,8 +245,11 @@ defmodule Module.Types.Apply do
 
         ## Map
         {:maps, :from_keys, [{[list(term()), term()], open_map()}]},
+        {:maps, :get, [{[term(), open_map()], term()}]},
         {:maps, :is_key, [{[term(), open_map()], boolean()}]},
         {:maps, :keys, [{[open_map()], dynamic(list(term()))}]},
+        {:maps, :take,
+         [{[term(), open_map()], tuple([term(), open_map()]) |> union(atom([:error]))}]},
         {:maps, :to_list, [{[open_map()], dynamic(list(tuple([term(), term()])))}]},
         {:maps, :values, [{[open_map()], dynamic(list(term()))}]}
       ] do
@@ -422,6 +425,43 @@ defmodule Module.Types.Apply do
     end
   end
 
+  defp remote_apply(:maps, :get, _info, [key, map] = args_types, stack) do
+    case map_get(map, key) do
+      {_, value} ->
+        {:ok, return(value, args_types, stack)}
+
+      :badmap ->
+        {:error, badremote(:maps, :get, 2)}
+
+      :error ->
+        {:error, {:badkeydomain, map, key, []}}
+    end
+  end
+
+  defp remote_apply(:maps, :take, _info, [key, map] = args_types, stack) do
+    case map_update(map, key, not_set()) do
+      # We could suggest to use :maps.delete if the key always exists
+      # but :maps.take/2 means calling Erlang directly, so we are fine.
+      {value, descr, errors} ->
+        result = tuple([value, descr])
+
+        result =
+          if errors == [] and not gradual?(map) do
+            result
+          else
+            union(result, atom([:error]))
+          end
+
+        {:ok, return(result, args_types, stack)}
+
+      :badmap ->
+        {:error, badremote(:maps, :take, 2)}
+
+      {:error, _errors} ->
+        {:ok, return(atom([:error]), args_types, stack)}
+    end
+  end
+
   defp remote_apply(:maps, :to_list, _info, [map], stack) do
     case map_to_list(map) do
       {:ok, list_type} -> {:ok, return(list_type, [map], stack)}
@@ -1033,6 +1073,58 @@ defmodule Module.Types.Apply do
     }
   end
 
+  def format_diagnostic({{:badkeydomain, map, key, errors}, _args_types, mfac, expr, context}) do
+    {mod, fun, arity, _converter} = mfac
+    traces = collect_traces(expr, context)
+    mfa_or_fa = if mod, do: Exception.format_mfa(mod, fun, arity), else: "#{fun}/#{arity}"
+
+    error_key =
+      Enum.reduce(errors, none(), fn
+        {:badkey, key}, acc -> union(atom([key]), acc)
+        {:baddomain, domain}, acc -> union(domain, acc)
+      end)
+
+    %{
+      details: %{typing_traces: traces},
+      message:
+        IO.iodata_to_binary([
+          """
+          incompatible types given to #{mfa_or_fa}:
+
+              #{expr_to_string(expr) |> indent(4)}
+
+          """,
+          if errors == [] do
+            """
+            the map:
+
+                #{to_quoted_string(map) |> indent(4)}
+
+            does not have the given keys:
+
+                #{to_quoted_string(key) |> indent(4)}
+
+            """
+          else
+            """
+            expected the map:
+
+                #{to_quoted_string(map) |> indent(4)}
+
+            to have the keys:
+
+                #{to_quoted_string(key) |> indent(4)}
+
+            but the following keys are missing:
+
+                #{to_quoted_string(error_key) |> indent(4)}
+            """
+          end,
+          format_traces(traces)
+        ])
+    }
+  end
+
   def format_diagnostic({{:badremote, {_, domain, clauses}}, args_types, mfac, expr, context}) do
     domain = domain(domain, clauses)
     {mod, fun, arity, converter} = mfac

lib/elixir/lib/module/types/descr.ex

@@ -2970,7 +2970,7 @@ defmodule Module.Types.Descr do
   The gradual aspect of `key_descr` does not impact the return type.
 
   It returns `{type, descr, errors}`, `:badmap`, `{:error, errors}`.
-  The list of `errors` may be empty, which implies the a bad domain.
+  The list of `errors` may be empty, which implies a bad domain.
   The `return_type?` flag is used for optimizations purposes. If set to false,
   the returned `type` should not be used, as it will be imprecise.
   """
@@ -3023,6 +3023,8 @@ defmodule Module.Types.Descr do
               empty?(type)
             end)
 
+          # We can exceptionally check for none() here because
+          # we already check for empty downstream
           if dynamic_descr == none() do
             {:error, static_errors ++ dynamic_errors}
           else

lib/elixir/test/elixir/module/types/map_test.exs

@@ -11,6 +11,97 @@ defmodule Module.Types.MapTest do
   import Module.Types.Descr
   defmacro domain_key(arg) when is_atom(arg), do: [arg]
 
+  describe ":maps.take/2" do
+    test "checking" do
+      assert typecheck!(:maps.take(:key, %{})) == atom([:error])
+
+      assert typecheck!(:maps.take(:key, %{key: 123})) |> equal?(tuple([integer(), empty_map()]))
+
+      assert typecheck!([x], :maps.take(:key, x))
+             |> equal?(
+               union(
+                 dynamic(tuple([term(), open_map(key: not_set())])),
+                 atom([:error])
+               )
+             )
+
+      assert typecheck!([condition?, x], :maps.take(if(condition?, do: :foo, else: :bar), x))
+             |> equal?(
+               union(
+                 dynamic(
+                   tuple([
+                     term(),
+                     union(
+                       open_map(foo: not_set()),
+                       open_map(bar: not_set())
+                     )
+                   ])
+                 ),
+                 atom([:error])
+               )
+             )
+
+      assert typecheck!([x], :maps.take(integer(), x))
+             |> equal?(
+               union(
+                 dynamic(tuple([term(), open_map()])),
+                 atom([:error])
+               )
+             )
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = :maps.take(:key, x)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!([x = []], :maps.take(:foo, x)) =~
+               "incompatible types given to :maps.take/2"
+    end
+  end
+
+  describe "Map.fetch!/2" do
+    test "errors" do
+      assert typeerror!(Map.fetch!(%{}, :foo)) =~
+               """
+               incompatible types given to Map.fetch!/2:
+
+                   Map.fetch!(%{}, :foo)
+
+               the map:
+
+                   empty_map()
+
+               does not have the given keys:
+
+                   :foo
+
+               """
+
+      assert typeerror!(Map.fetch!(%{}, 123)) =~
+               """
+               incompatible types given to Map.fetch!/2:
+
+                   Map.fetch!(%{}, 123)
+
+               the map:
+
+                   empty_map()
+
+               does not have the given keys:
+
+                   integer()
+
+               """
+    end
+  end
+
   describe "Map.to_list/1" do
     test "checking" do
       assert typecheck!([x = %{}], Map.to_list(x)) == dynamic(list(tuple([term(), term()])))