Support do ... rescue ... end blocks. Fixes gyson#34.

lib/ex_type/checker.ex

@@ -22,6 +22,15 @@ defmodule ExType.Checker do
 
   @spec eval(Context.t(), any()) :: {Context.t(), Type.t()}
 
+  def eval(context, {:do, block}) do
+    eval(context, block)
+  end
+
+  # handle rescue ... end
+  def eval(context, {:rescue, args}) do
+    block_of_bindings(context, args)
+  end
+
   def eval(context, {:%{}, _, []}) do
     {context, %Type.Map{key: %Type.Any{}, value: %Type.Any{}}}
   end
@@ -439,25 +448,7 @@ defmodule ExType.Checker do
 
   # handle receive do ... end
   def eval(context, {:receive, _, [[do: args]]}) do
-    t =
-      for {:->, _, [[left], right]} <- args do
-        new_context =
-          case left do
-            {:when, _, [when_left, when_right]} ->
-              context
-              |> Unification.unify_pattern(when_left, %Type.Any{})
-              |> Unification.unify_guard(when_right)
-
-            _ ->
-              Unification.unify_pattern(context, left, %Type.Any{})
-          end
-
-        {_, type} = eval(new_context, right)
-        type
-      end
-      |> Typespec.union_types()
-
-    {context, t}
+    block_of_bindings(context, args)
   end
 
   # function call, e.g. 1 + 1
@@ -574,6 +565,27 @@ defmodule ExType.Checker do
     )
   end
 
+  defp block_of_bindings(context, args) do
+    t = for {:->, _, [[left], right]} <- args do
+        new_context =
+          case left do
+            {:when, _, [when_left, when_right]} ->
+              context
+              |> Unification.unify_pattern(when_left, %Type.Any{})
+              |> Unification.unify_guard(when_right)
+
+            _ ->
+              Unification.unify_pattern(context, left, %Type.Any{})
+          end
+
+        {_, type} = eval(new_context, right)
+        type
+      end
+      |> Typespec.union_types()
+
+    {context, t}
+  end
+
   def exact_pattern_match(atom, %Type.Atom{literal: true, value: atom}) when is_atom(atom) do
     true
   end

lib/ex_type/custom_env.ex

@@ -104,12 +104,13 @@ defmodule ExType.CustomEnv do
     end
   end
 
-  def save_def(module, call, do: block) do
+  # Allow `do ... end` as well as `do ... rescue ... end`
+  def save_def(module, call, [{:do, _} | _] = block) do
     Module.register_attribute(module, :ex_type_def, accumulate: true, persist: true)
     Module.put_attribute(module, :ex_type_def, {call, block})
   end
 
-  def save_defp(module, call, do: block) do
+  def save_defp(module, call, [{:do, _} | _] = block) do
     Module.register_attribute(module, :ex_type_defp, accumulate: true, persist: true)
     Module.put_attribute(module, :ex_type_defp, {call, block})
   end

lib/ex_type/parser.ex

@@ -70,28 +70,36 @@ defmodule ExType.Parser do
 
     expanded_body =
       block
-      |> Macro.postwalk(fn
-        # support T.inspect
-        {:., m1, [{:__aliases__, m2, [:T]}, :inspect]} ->
-          {:., m1, [{:__aliases__, m2, [:ExType, :T, :TypeCheck]}, :inspect]}
-
-        # support T.assert
-        {{:., m1, [{:__aliases__, m2, [:T]}, :assert]}, m3, [arg]} ->
-          case arg do
-            {operator, _, [left, right]} when operator in [:==, :"::", :<, :>] ->
-              {{:., m1, [{:__aliases__, m2, [:ExType, :T, :TypeCheck]}, :assert]}, m3,
-               [operator, left, Macro.escape(right)]}
-          end
-
-        code ->
-          code
+      |> Enum.map(fn {block_kind, block_part} ->
+          {block_kind,
+            block_part
+              |> Macro.postwalk(&walk_support_inspect_assert/1)
+              |> replace_module_macro(module_name)
+              |> expand_all(updated_env)
+          }
       end)
-      |> replace_module_macro(module_name)
-      |> expand_all(updated_env)
 
     {name, expanded_args, expanded_guard, expanded_body}
   end
 
+  # support T.inspect
+  defp walk_support_inspect_assert({:., m1, [{:__aliases__, m2, [:T]}, :inspect]}) do
+    {:., m1, [{:__aliases__, m2, [:ExType, :T, :TypeCheck]}, :inspect]}
+  end
+
+  # support T.assert
+  defp walk_support_inspect_assert({{:., m1, [{:__aliases__, m2, [:T]}, :assert]}, m3, [arg]}) do
+    case arg do
+      {operator, _, [left, right]} when operator in [:==, :"::", :<, :>] ->
+        {{:., m1, [{:__aliases__, m2, [:ExType, :T, :TypeCheck]}, :assert]}, m3,
+          [operator, left, Macro.escape(right)]}
+    end
+  end
+
+  defp walk_support_inspect_assert(code) do
+    code
+  end
+
   # replace __MODULLE__ with actual module name, removed "ExType.Module" prefix
   # TODO: handle this properly when replace :elixir_expand.expand, especially for nested case
   defp replace_module_macro(block, module_name) do

lib/ex_type/typespec.ex

@@ -792,9 +792,19 @@ defmodule ExType.Typespec do
           end
           |> case do
             {:ok, new_fn_context} ->
-              {_, result_type} = ExType.Checker.eval(new_fn_context, body)
-
-              [result_type]
+              case body do
+                [do: block] ->
+                  {_context, type} = ExType.Checker.eval(new_fn_context, {:do, block})
+                  [type]
+                # For `do ... rescue ... end`-blocks, the result type can either be the first or the second part
+                [do: block_try, rescue: block_rescue] ->
+                  {_context, type_try} = ExType.Checker.eval(new_fn_context, {:do, block_try})
+                  {_context, type_rescue} = ExType.Checker.eval(new_fn_context, {:rescue, block_rescue})
+                  [Type.union([type_try, type_rescue])]
+                block ->
+                  {_context, type} = ExType.Checker.eval(new_fn_context, block)
+                  [type]
+              end
 
             :not_match ->
               []