Support when in typespec

Author: Premwoik

lib/gradient/ast_specifier.ex

@@ -478,6 +478,24 @@ defmodule Gradient.AstSpecifier do
     |> pass_tokens(tokens)
   end
 
+  def spec_mapper({:type, anno, :constraint, [subtype, vt]}, tokens, opts) do
+    {:ok, _, anno, opts, _} = get_line(anno, opts)
+    {subtype, _} = spec_mapper(subtype, tokens, opts)
+    vt = context_mapper_map(vt, tokens, opts, &spec_mapper/3)
+
+    {:type, anno, :constraint, [subtype, vt]}
+    |> pass_tokens(tokens)
+  end
+
+  def spec_mapper({:type, anno, :bounded_fun, [fn_type, when_type]}, tokens, opts) do
+    {:ok, _line, anno, opts, _} = get_line(anno, opts)
+    {fn_type, _} = spec_mapper(fn_type, tokens, opts)
+    when_type = context_mapper_map(when_type, tokens, opts, &spec_mapper/3)
+
+    {:type, anno, :bounded_fun, [fn_type, when_type]}
+    |> pass_tokens(tokens)
+  end
+
   def spec_mapper({:type, anno, type_name, args}, tokens, opts) do
     {:ok, _line, anno, opts, _} = get_line(anno, opts)
     new_args = context_mapper_map(args, tokens, opts, &spec_mapper/3)

test/examples/Elixir.TypespecWhen.beam

@@ -0,0 +1,12 @@
+defmodule TypespecWhen do
+  @spec foo({:a, x}) :: {:a, x} | {:b, x} when x: term()
+  def foo({:a, x}) do
+    case x do
+      :foo ->
+        {:a, x}
+
+      _ ->
+        {:b, x}
+    end
+  end
+end

test/examples/typespec_when.ex

@@ -1294,6 +1294,36 @@ defmodule Gradient.AstSpecifierTest do
             ]} = recv
   end
 
+  test "typespec when" do
+    {tokens, ast} = load("/Elixir.TypespecWhen.beam", "/typespec_when.ex")
+
+    [spec | _] =
+      AstSpecifier.run_mappers(ast, tokens)
+      |> filter_specs()
+      |> Enum.reverse()
+
+    assert {:attribute, 2, :spec,
+            {{:foo, 1},
+             [
+               {:type, 2, :bounded_fun,
+                [
+                  {:type, 2, :fun,
+                   [
+                     {:type, 2, :product, [{:type, 2, :tuple, [{:atom, 2, :a}, {:var, 2, :x}]}]},
+                     {:type, 2, :union,
+                      [
+                        {:type, 2, :tuple, [{:atom, 2, :a}, {:var, 2, :x}]},
+                        {:type, 2, :tuple, [{:atom, 2, :b}, {:var, 2, :x}]}
+                      ]}
+                   ]},
+                  [
+                    {:type, 2, :constraint,
+                     [{:atom, 2, :is_subtype}, [{:var, 2, :x}, {:type, 2, :term, []}]]}
+                  ]
+                ]}
+             ]}} = spec
+  end
+
   test "typespec" do
     {tokens, ast} = load("/Elixir.Typespec.beam", "/typespec.ex")