Add function signature utilities

src/ExpressionUtils.jl

@@ -3,6 +3,7 @@ module ExpressionUtils
 using Compat
 
 export is_funcdef_expr, get_funcdef_expr, funcdef_longform
+export funcdef_name, funcdef_params, funcdef_args, funcdef_argnames, funcdef_argtypeexprs
 export walk, expr_replace, expr_bind, expr_bindings
 
 include("blocks.jl")

src/functions.jl

@@ -120,3 +120,97 @@ function funcdef_longform(ex::Expr)
     end
     throw(ArgumentError(string("expected a function definition, got ", ex)))
 end
+
+"""
+    funcdef_name(sig)
+
+Return the function's Symbol given a signature expression `sig`. `sig`
+is the first argument of the function definition expression.
+"""
+function funcdef_name(sig::Expr)
+    sig.head == :call || throw(ArgumentError(string("expected call expression, got ", sig)))
+    decl = sig.args[1]
+    isa(decl, Symbol) && return decl::Symbol
+    if isa(decl, Expr) && (decl::Expr).head == :curly
+        return ((decl::Expr).args[1])::Symbol
+    end
+    throw(ArgumentError(string("unexpected call expression ", sig)))
+end
+
+"""
+    funcdef_params(sig)
+
+Return a Vector{Any} of the function's parameters (the part inside the
+curlies) given a signature expression `sig`. `sig` is the first
+argument of the function definition expression.
+
+The elements of the return may be Symbols or Exprs, where the latter
+is used for expressions like `T<:Integer`.
+"""
+function funcdef_params(sig::Expr)
+    sig.head == :call || throw(ArgumentError(string("expected call expression, got ", sig)))
+    ret = []
+    decl = sig.args[1]
+    isa(decl, Symbol) && return ret
+    if isa(decl, Expr) && (decl::Expr).head == :curly
+        for i = 2:length((decl::Expr).args)
+            push!(ret, (decl::Expr).args[i])
+        end
+        return ret
+    end
+    throw(ArgumentError(string("unexpected call expression ", sig)))
+end
+
+"""
+    funcdef_args(sig)
+
+Return a Vector{Any} of the function's argument expressions, which
+includes both the variable name and any type declaration, if any.
+
+The elements of the return may be Symbols or Exprs, where the latter
+is used for expressions like `x::Integer`.
+
+See also `funcdef_argnames` and `funcdef_argtypes`.
+"""
+function funcdef_args(sig::Expr)
+    sig.head == :call || throw(ArgumentError(string("expected call expression, got ", sig)))
+    return sig.args[2:end]
+end
+
+"""
+    funcdef_argnames(sig)
+
+Return a vector of the function's argument names.  These are typically
+Symbols, except for non-canonical syntax like that of Traitor.jl.
+
+See also `funcdef_args` and `funcdef_argtypeexprs`.
+"""
+function funcdef_argnames(sig::Expr)
+    sig.head == :call || throw(ArgumentError(string("expected call expression, got ", sig)))
+    return [argname(a) for a in sig.args[2:end]]
+end
+
+argname(s::Symbol) = s
+function argname(ex::Expr)
+    ex.head == :(::) || throw(ArgumentError("expected :(::) expression, got ", ex))
+    ex.args[1]
+end
+
+"""
+    funcdef_argtypeexprs(sig)
+
+Return a vector of the function's argument type expressions.  These
+may be Symbols (`:Int`) or expressions (`AbstractArray{T}`).
+
+See also `funcdef_args` and `funcdef_argnames`.
+"""
+function funcdef_argtypeexprs(sig::Expr)
+    sig.head == :call || throw(ArgumentError(string("expected call expression, got ", sig)))
+    return [argtypeexpr(a) for a in sig.args[2:end]]
+end
+
+argtypeexpr(s::Symbol) = :Any
+function argtypeexpr(ex::Expr)
+    ex.head == :(::) || throw(ArgumentError("expected :(::) expression, got ", ex))
+    ex.args[2]
+end

test/functions.jl

@@ -1,43 +1,70 @@
 @testset "Functions" begin
-    ex1 = :(f(x) = 1)
-    ex2 = :(@inline f(x) = 1)
-    ex3 = :(Base.@pure @inline f(x) = 1)
-    ex4 = quote
-        f(x) = 1
-    end
-    ex5 = :(function f(x) 1 end)
-    ex6 = quote
-        function f(x)
-            1
+
+    @testset "funcdef" begin
+        ex1 = :(f(x) = 1)
+        ex2 = :(@inline f(x) = 1)
+        ex3 = :(Base.@pure @inline f(x) = 1)
+        ex4 = quote
+            f(x) = 1
         end
-    end
-    ex7 = quote
-        @inline function f(x)
-            1
+        ex5 = :(function f(x) 1 end)
+        ex6 = quote
+            function f(x)
+                1
+            end
         end
-    end
-    ex8 = :(immutable NotAFunction end)
+        ex7 = quote
+            @inline function f(x)
+                1
+            end
+        end
+        ex8 = :(immutable NotAFunction end)
 
-    for ex in (ex1, ex5)
-        @test is_funcdef_expr(ex)
-    end
-    for ex in (ex2, ex3, ex4, ex6, ex7, ex8)
-        @test !is_funcdef_expr(ex)
+        for ex in (ex1, ex5)
+            @test is_funcdef_expr(ex)
+        end
+        for ex in (ex2, ex3, ex4, ex6, ex7, ex8)
+            @test !is_funcdef_expr(ex)
+        end
+        for (ex,contains) in ((ex1,false), (ex2,true), (ex3,true), (ex4,true),
+                              (ex5,false), (ex6,true), (ex7,true))
+            def, c, container = get_funcdef_expr(ex)
+            @test c == contains
+            @test is_funcdef_expr(def)
+            @test isa(container, Expr)
+            if contains
+                @test container.args[end] === def
+            end
+            lf = funcdef_longform(def)
+            @test lf.head == :function
+        end
+        @test_throws ArgumentError get_funcdef_expr(ex8)
+        def, c, container = get_funcdef_expr(ex8, false)
+        @test isa(def, Expr)
+        @test !is_funcdef_expr(def)
     end
-    for (ex,contains) in ((ex1,false), (ex2,true), (ex3,true), (ex4,true),
-                          (ex5,false), (ex6,true), (ex7,true))
-        def, c, container = get_funcdef_expr(ex)
-        @test c == contains
-        @test is_funcdef_expr(def)
-        @test isa(container, Expr)
-        if contains
-            @test container.args[end] === def
+
+    @testset "signature" begin
+        # These are all canonicalized, which is OK because we've tested funcdef_longform
+        ex1 = :(function f(x) 1 end)
+        ex2 = :(function f(x::Int) 1 end)
+        ex3 = :(function f{T}(x::T) 1 end)
+        ex4 = :(function f{R,S}(x::R, y::S, z) 1 end)
+        ex5 = :(function f{T<:Integer}(x::T, y::String) 1 end)
+        ex6 = :(function f{T}(x::AbstractArray{T}) 1 end)
+        for (ex, params, args, names, types) in ((ex1, Symbol[], [:x], [:x], [:Any]),
+                                                 (ex2, Symbol[], [:(x::Int)], [:x], [:Int]),
+                                                 (ex3, [:T], [:(x::T)], [:x], [:T]),
+                                                 (ex4, [:R,:S], [:(x::R), :(y::S), :z], [:x, :y, :z], [:R, :S, :Any]),
+                                                 (ex5, [:(T<:Integer)], [:(x::T), :(y::String)], [:x, :y], [:T, :String]),
+                                                 (ex6, [:T], [:(x::AbstractArray{T})], [:x], [:(AbstractArray{T})]))
+            sig = ex.args[1]
+            @test funcdef_name(sig) == :f
+            @test funcdef_params(sig) == params
+            @test funcdef_args(sig) == args
+            @test funcdef_argnames(sig) == names
+            @test funcdef_argtypeexprs(sig) == types
         end
-        lf = funcdef_longform(def)
-        @test lf.head == :function
     end
-    @test_throws ArgumentError get_funcdef_expr(ex8)
-    def, c, container = get_funcdef_expr(ex8, false)
-    @test isa(def, Expr)
-    @test !is_funcdef_expr(def)
+
 end