Add support for default values to TYPEDSIGNATURES

Project.toml

@@ -6,12 +6,16 @@ version = "0.9.3"
 LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
 [compat]
+CodeTracking = "1.3.6"
+REPL = "1"
 julia = "1"
 
 [extras]
+CodeTracking = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Markdown", "Pkg", "Test"]
+test = ["Markdown", "Pkg", "Test", "REPL", "CodeTracking"]

src/DocStringExtensions.jl

@@ -86,6 +86,7 @@ export interpolation
 
 # Includes.
 
+include("parsing.jl")
 include("utilities.jl")
 include("abbreviations.jl")
 include("templates.jl")

src/abbreviations.jl

@@ -340,7 +340,11 @@ The singleton type for [`TYPEDSIGNATURES`](@ref) abbreviations.
 
 $(:FIELDS)
 """
-struct TypedMethodSignatures <: Abbreviation end
+struct TypedMethodSignatures <: Abbreviation
+    expr::Union{Nothing, Expr}
+end
+
+interpolation(::TypedMethodSignatures, expr) = TypedMethodSignatures(expr)
 
 """
 An [`Abbreviation`](@ref) for including a simplified representation of all the method
@@ -358,21 +362,24 @@ f(x::Int, y::Int; a, b...)
 ```
 ````
 """
-const TYPEDSIGNATURES = TypedMethodSignatures()
+const TYPEDSIGNATURES = TypedMethodSignatures(nothing)
+
+function format(x::TypedMethodSignatures, buf, doc)
+    binding = doc.data[:binding]
+    typesig = doc.data[:typesig]
+    modname = doc.data[:module]
+    func = Docs.resolve(binding)
 
-function format(::TypedMethodSignatures, buf, doc)
-    local binding = doc.data[:binding]
-    local typesig = doc.data[:typesig]
-    local modname = doc.data[:module]
-    local func = Docs.resolve(binding)
     # TODO: why is methodgroups returning invalid methods?
     # the methodgroups always appears to return a Vector and the size depends on whether parametric types are used
     # and whether default arguments are used
-    local groups = methodgroups(func, typesig, modname)
+    groups = methodgroups(func, typesig, modname)
     if !isempty(groups)
         group = groups[end]
+        ast_info = parse_call(x.expr)
         println(buf)
         println(buf, "```julia")
+
         for (i, method) in enumerate(group)
             N = length(arguments(method))
             # return a list of tuples that represent type signatures
@@ -395,9 +402,11 @@ function format(::TypedMethodSignatures, buf, doc)
             else
                 t = tuples[findfirst(f, tuples)]
             end
-            printmethod(buf, binding, func, method, t)
+
+            printmethod(buf, binding, func, ast_info.args, ast_info.kwargs, t)
             println(buf)
         end
+
         println(buf, "\n```\n")
     end
 end

src/parsing.jl

@@ -0,0 +1,122 @@
+Base.@kwdef struct ASTArg
+    name::Union{Symbol, Nothing} = nothing
+    type = nothing
+    default = nothing
+    variadic::Bool = false
+end
+
+# Parse an argument with a type annotation.
+# Example input: `x::Int`
+function parse_arg_with_type(arg_expr::Expr)
+    if arg_expr.head != :(::)
+        throw(ArgumentError("Argument is not a :(::) expr"))
+    end
+
+    n_expr_args = length(arg_expr.args)
+    return if n_expr_args == 1
+        # '::Int'
+        ASTArg(; type=arg_expr.args[1])
+    elseif n_expr_args == 2
+        # 'x::Int'
+        ASTArg(; name=arg_expr.args[1], type=arg_expr.args[2])
+    end
+end
+
+# Parse an argument with a default value.
+# Example input: `x=5`
+function parse_arg_with_default(arg_expr::Expr)
+    if arg_expr.head != :kw
+        throw(ArgumentError("Argument is not a :kw expr"))
+    end
+
+    if arg_expr.args[1] isa Symbol
+        # This is an argument without a type annotation
+        ASTArg(; name=arg_expr.args[1], default=arg_expr.args[2])
+    else
+        # This is an argument with a type annotation
+        tmp = parse_arg_with_type(arg_expr.args[1])
+        ASTArg(; name=tmp.name, type=tmp.type, default=arg_expr.args[2])
+    end
+end
+
+# Parse a list of expressions, assuming the list is an argument list containing
+# positional/keyword arguments.
+# Example input: `(x, y::Int; z=5, kwargs...)`
+function parse_arglist!(exprs, args, kwargs, is_kwarg_list=false)
+    list = is_kwarg_list ? kwargs : args
+
+    for arg_expr in exprs
+        if arg_expr isa Symbol
+            # Plain argument name with no type or default value
+            push!(list, ASTArg(; name=arg_expr))
+        elseif arg_expr.head == :(::)
+            # With a type annotation
+            push!(list, parse_arg_with_type(arg_expr))
+        elseif arg_expr.head == :kw
+            # With a default value (and possibly a type annotation)
+            push!(list, parse_arg_with_default(arg_expr))
+        elseif arg_expr.head == :parameters
+            # Keyword arguments
+            parse_arglist!(arg_expr.args, args, kwargs, true)
+        elseif arg_expr.head === :...
+            # Variadic argument
+            if arg_expr.args[1] isa Symbol
+                # Without a type annotation
+                push!(list, ASTArg(; name=arg_expr.args[1], variadic=true))
+            elseif arg_expr.args[1].head === :(::)
+                # With a type annotation
+                arg_expr = arg_expr.args[1]
+                push!(list, ASTArg(; name=arg_expr.args[1], type=arg_expr.args[2], variadic=true))
+            else
+                Meta.dump(arg_expr)
+                error("Couldn't parse variadic Expr in arg list (printed above)")
+            end
+        else
+            Meta.dump(arg_expr)
+            error("Couldn't parse Expr in arg list (printed above)")
+        end
+    end
+end
+
+# Find a :call expression within an Expr. This will take care of ignoring other
+# tokens like `where` clauses.
+function find_call_expr(expr::Expr)
+    if expr.head === :macrocall && expr.args[1] === Symbol("@generated")
+        # If this is a generated function, find the first := expr to find
+        # the :call expr.
+        assignment_idx = findfirst(x -> x isa Expr && x.head === :(=), expr.args)
+
+        expr.args[assignment_idx].args[1]
+    elseif expr.head === :(=)
+        find_call_expr(expr.args[1])
+    elseif expr.head == :where
+        # Function with one or more `where` clauses
+        find_call_expr(expr.args[1])
+    elseif expr.head === :function
+        find_call_expr(expr.args[1])
+    elseif expr.head === :call
+        expr
+    else
+        Meta.dump(expr)
+        error("Can't parse current expr (printed above)")
+    end
+end
+
+# Parse an expression to find a :call expr, and return as much information as
+# possible about the arguments.
+# Example input: `foo(x) = x^2`
+function parse_call(expr::Expr)
+    Base.remove_linenums!(expr)
+    expr = find_call_expr(expr)
+
+    if expr.head != :call
+        throw(ArgumentError("Argument is not a :call, cannot parse it."))
+    end
+
+    args = ASTArg[]
+    kwargs = ASTArg[]
+    # Skip the first argument because that's just the function name
+    parse_arglist!(expr.args[2:end], args, kwargs)
+
+    return (; args, kwargs)
+end

src/utilities.jl

@@ -321,28 +321,7 @@ function find_tuples(typesig)
     end
 end
 
-"""
-$(:TYPEDSIGNATURES)
-
-Print a simplified representation of a method signature to `buffer`. Some of these
-simplifications include:
-
-  * no `TypeVar`s;
-  * no types;
-  * no keyword default values;
-
-# Examples
-
-```julia
-f(x::Int; a = 1, b...) = x
-sig = printmethod(Docs.Binding(Main, :f), f, first(methods(f)))
-```
-"""
-function printmethod(buffer::IOBuffer, binding::Docs.Binding, func, method::Method, typesig)
-    # TODO: print qualified?
-    local args = string.(arguments(method))
-    local kws = string.(keywords(func, method))
-
+function format_args(args::Vector{ASTArg}, typesig)
     # find inner tuple type
     function find_inner_tuple_type(t)
         # t is always either a UnionAll which represents a generic type or a Tuple where each parameter is the argument
@@ -383,22 +362,65 @@ function printmethod(buffer::IOBuffer, binding::Docs.Binding, func, method::Meth
             collect(typesig.types)
 
     args = map(args, argtypes) do arg,t
+        name = ""
         type = ""
         suffix = ""
+        default_value = ""
+
+        if !isnothing(arg.name)
+            name = arg.name
+        end
         if isvarargtype(t)
             t = vararg_eltype(t)
             suffix = "..."
+        elseif arg.variadic
+            # This extra branch is here for kwargs, where we don't have type
+            # information.
+            suffix = "..."
         end
-        if t!==Any
+        if t !== Any
             type = "::$t"
         end
+        if !isnothing(arg.default)
+            default_value = "=$(arg.default)"
+        end
 
-        "$arg$type$suffix"
+        "$name$type$suffix$default_value"
     end
 
+    return args
+end
+
+"""
+$(:TYPEDSIGNATURES)
+
+Print a simplified representation of a method signature to `buffer`. Some of these
+simplifications include:
+
+  * no `TypeVar`s;
+  * no types;
+  * no keyword default values;
+
+# Examples
+
+```julia
+f(x::Int; a = 1, b...) = x
+sig = printmethod(Docs.Binding(Main, :f), f, first(methods(f)))
+```
+"""
+function printmethod(buffer::IOBuffer, binding::Docs.Binding, func,
+                     args::Vector{ASTArg}, kws::Vector{ASTArg}, typesig)
+    formatted_args = format_args(args, typesig)
+    # We don't have proper type information for keyword arguments like we do
+    # with `typesig` for positional arguments, so we assume they're all Any. An
+    # alternative would be to use the types extracted from the AST, but that
+    # might not exactly match the types of positional arguments (e.g. an alias
+    # type would be printed as the underlying type for positional arguments but
+    # under the alias for keyword arguments).
+    formatted_kws = format_args(kws, NTuple{length(kws), Any})
     rt = Base.return_types(func, typesig)
 
-    return printmethod_format(buffer, string(binding.var), args, string.(kws);
+    return printmethod_format(buffer, string(binding.var), formatted_args, formatted_kws;
         return_type =
             length(rt) >= 1 && rt[1] !== Nothing && rt[1] !== Union{} ?
             " -> $(rt[1])" : "")

test/runtests.jl

@@ -1,4 +1,6 @@
 using DocStringExtensions
+import DocStringExtensions: TypedMethodSignatures
+
 using Test
 import Markdown
 import LibGit2