Can typespec be added to the function like comments are?

[ivan-srsen]

No description provided.

[doorgan]

Yes, and a good way to do it is by using Zippers:

As always with AST, it can be a bit tricky:

code = """
defmodule Hello do
  def foo do
    :ok
  end
end
"""

ast = Sourceror.parse_string!(code)

{:defmodule,
 [
   trailing_comments: [],
   leading_comments: [],
   do: [line: 1, column: 17],
   end: [line: 5, column: 1],
   line: 1,
   column: 1
 ],
 [
   {:__aliases__,
    [trailing_comments: [], leading_comments: [], last: [line: 1, column: 11], line: 1, column: 11],
    [:Hello]},
   [
     {{:__block__, [trailing_comments: [], leading_comments: [], line: 1, column: 17], [:do]},
      {:def,
       [
         trailing_comments: [],
         leading_comments: [],
         do: [line: 2, column: 11],
         end: [line: 4, column: 3],
         line: 2,
         column: 3
       ],
       [
         {:foo, [trailing_comments: [], leading_comments: [], line: 2, column: 7], nil},
         [
           {{:__block__, [trailing_comments: [], leading_comments: [], line: 2, column: 11], [:do]},
            {:__block__, [trailing_comments: [], leading_comments: [], line: 3, column: 5], [:ok]}}
         ]
       ]}}
   ]
 ]}

alias Sourceror.Zipper

modified =
  ast
  |> Zipper.zip()
  |> Zipper.traverse(fn
    {{:def, _, [{:foo, _, _}, _]}, _path} = zipper ->
      spec =
        quote do
          @spec foo() :: :ok
        end

      case Zipper.up(zipper) do
        {{{:__block__, _, _} = do_block, content}, path} ->
          # This is the only expression inside the `do` block, so we need to wrap it in a block
          # We're at {do_block, content}
          {{do_block, {:__block__, [], [content]}}, path}
          # Go down to do_block
          |> Zipper.down()
          # Go right to [content]
          |> Zipper.right()
          # We're in [content], which is a list, so we go down inside of it
          |> Zipper.down()
          # Finally, we can insert the spec
          |> Zipper.insert_left(spec)

        _ ->
          # We're inside a block, so we can just insert the spec
          Zipper.insert_left(zipper, spec)
      end

    other ->
      other
  end)
  |> Zipper.node()
  |> Sourceror.to_string()
  |> IO.puts()

defmodule Hello do
  @spec foo() :: :ok
  def foo do
    :ok
  end
end

:ok

The hassle here is that the AST is represented differently depending on if the do block has one or more expressions inside of it. Maybe we could be a bit smarter with Sourceror and make it be always teh same, but at the moment we carry over the quirks of the elixir parser.

I encourage you to parse the AST you're interested in and experiment a bit with it, as most of the time you want to recognize a pattern, or a set of patterns, and match against them to manipulate them.

There's also room for improvement for instance to avoid having to do the down |> right |> down dance, there's an open issue to expand Sourceror's api to make navigation easier in #64