Adding static type checking to Julia in 100 lines of code

4 years ago

source link: https://www.tuicool.com/articles/IZ3I3qM
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

While inJuliatype annotations are optional, the type system is designed in such a way that the types of variables in a function can be inferred for given inputs. By taking advantage of this we can add checks before a function ever gets called that the return types are as expected.

Given that Julia is dynamic (types can be added at any time) and generic by default (meaning a huge number of types would need to be checked in many cases) this kind of checks might not be very useful in practice. However, it's a good excuse to explore Julia's meta-programming and introspection tools.

We will first deal with the simplest case possible (a method annotated with concrete types) and then generalize to parametric methods.

Non-parametric methods

We will deal with the simplest case to begin with. We want to build a macro @checked that will take a return type and a function definition as input and add a test that will check that the inferred return type matches the declared one after the function definition. In short, transforms the code:

@checked Int16 function add(x::Int8,
 y::Int16)
    
\n
  x+y\nend

Julia

Into:

function add(x::Int8,
 y::Int16)
    
\n
  x+y\nend\n@assert inferred(add,
 Tuple{Int8,
 Int16}
)
 == Int16

Julia

In order to do this we need to:

Extract the type annotation ( Int8 , Int16 ) from the Abstract Syntax Tree ( AST )
Write an inferred method that return the inferred return type
Generate the final expression

For 1. we will use the MLStyle package which adds pattern matching for ASTs, making this task very easy.

For example if we want to extract the body of a comprehension [i^2 for i=1:10] using MLStyle's @match macro, we simply need to replace each element by wildcards [$body for $var=$range] :

using MLStyle\n
\nmatch_signature(ast)
 = @match ast begin\n
    quote\n
        $(::LineNumberNode)
\n
        
[$body for $var=$range]
\n
    end => body\nend\n
\nmatch_signature(
 quote 
[i^2 for i=1:10]
 end 
)

18.0s

MLStyle (Julia)

:(i ^ 2)

Note that Julia adds line numbers in its AST, so we also need to include a line number node in the pattern.

We can match our function signature in a similar fashion, returning the matched elements as a named tuple:

match_signature(ast)
 = @match ast begin\n
    quote\n
        $(::LineNumberNode)
\n
        function $name($(args...)
)
;
 $(block...)
 end\n
    end =>\n
    
(name = name,
\n
     args = args)
\nend\n
\nsig = match_signature(quote 
\n
    function add(x::Int8,y::Int16)
\n
        x+y\n
    end\nend)

1.7s

MLStyle (Julia)

(name = :add, args = Any[:(x::Int8), :(y::Int16)])

Similarly we define a method to extract type annotations ( x::T ):

match_type_annotation(ast)
 = @match ast begin :($(arg)::$(ann)
)
 => ann end\nmatch_type_annotation.(sig.args)

1.5s

MLStyle (Julia)

2-element Array{Symbol,1}:\n :Int8 \n :Int16

We can simply run type inference using Base.code_typed :

function inferred(f,args)
\n
    codeinfo = Base.code_typed(f,args)
[1]
\n
    codeinfo.second\nend\n
\n@assert inferred(
(x,y)
 -> x+y,
 Tuple{Int8,Int16}
)
 == Int16

0.6s

MLStyle (Julia)

When adding two integers of different type, the result gets automaticallypromoted(here to Int16 ). The promoted type can be computed using promote_type(T,K) .

We can now write a macro that will perform type checking:

macro checked(out,f)
\n
    sig  = match_signature(quote $f end)
\n
    args = match_type_annotation.(sig.args)
\n
  
\n
  \tchecks = quote\n
        @assert inferred($(sig.name)
,Tuple{$(args...)
}
)
 == $out\n
    end\n
    
\n
    esc(quote\n
        $f\n
        $checks\n
    end)
\nend\n
\n@macroexpand @checked Int16 function add(x::Int8,y::Int16)
\n
    x+y\nend

0.8s

MLStyle (Julia)

quote\n #= cell:10 =#\n function add(x::Int8, y::Int16)\n #= cell:16 =#\n x + y\n end\n #= cell:11 =#\n begin\n #= cell:6 =#\n if inferred(add, Tuple{Int8, Int16}) == Int16\n nothing\n else\n (Base.throw)((Base.AssertionError)("inferred(add, Tuple{Int8, Int16}) == Int16"))\n end\n end\nend

Here we use @macroexpand to see the code built by the macro. It first defines the function itself, and then adds a type check with the declared input and output types. We can now test that our macro work correctly:

@checked Int16 function add(x::Int8,y::Int16)
\n
    x+y\nend

0.4s

MLStyle (Julia)

If the inferred type is incorrect the macro will throw an AssertionError when the function is defined:

using Test\n@test_throws AssertionError @checked Int8 function add(x::Int8,y::Int16)
\n
    x+y\nend

1.0s

MLStyle (Julia)

Test Passed\n Thrown: AssertionError

Parametric methods

Next we will consider a simple parametric case, we will also replace the output type annotation by a predicate on the inferred type I to allow for more general (and useful) checks. For example in the following, the inferred type will be compared to the type of the first argument:

@checked I -> I == T function add(x::T,
 y::K)
 where 
{T<:Integer,
 K<:Integer}
 
\n
    x + y\nend

Julia

We first generalise our match_signature method to handle the parametric case, using an OR clause to handle functions without where :

match_signature(ast)
 = @match ast begin\n
    quote\n
        $(::LineNumberNode)
\n
        function $name($(args...)
)
 where 
{$(cov...)
}
;
 $(block...)
 end\n
    end ||\n
    quote\n
        $(::LineNumberNode)
\n
        function $name($(args...)
)
;
 $(block...)
 end\n
    end && Do(cov = 
[
]
)
 =>\n
    
(name = name,
\n
     args = args,
\n
     cov  = cov)
\nend\n
\nsig = match_signature(quote 
\n
    function add(x::T,y::K)
 where 
{T<:Integer,
 K<:Integer}
\n
        x+y\n
    end\nend)

2.1s

MLStyle (Julia)

(name = :add, args = Any[:(x::T), :(y::K)], cov = Any[:(T <: Integer), :(K <: Integer)])

We need a small method to extract subtype relations ( T <: Integer ):

match_subtype(ast)
 = @match ast begin :($(sub)<:$(sup)
)
 => 
(sub=sub,sup=sup)
 end\nmatch_subtype.(sig.cov)

1.5s

MLStyle (Julia)

2-element Array{NamedTuple{(:sub, :sup),Tuple{Symbol,Symbol}},1}:\n (sub = :T, sup = :Integer)\n (sub = :K, sup = :Integer)

In order to test every combination of parameters under the constrains above we need to enumerate all concret subtypes of Integer . This can be done by recursively calling InteractiveUtils.subtypes :

import InteractiveUtils.subtypes\n
\nfunction allsubtypes(T::Type)
\n
    isconcretetype(T)
 && return 
[T]
\n
    t = Type[
]
\n
    map(K->allsubtypes(K,t)
,
 subtypes(T)
)
\n
    t\nend\nallsubtypes(T,t)
 = isconcretetype(T)
 ? push!(t,T)
 : map(K->allsubtypes(K,t)
,
 subtypes(T)
)
\nallsubtypes(T::Symbol)
 = allsubtypes(Core.eval(Main,T)
)#convert Symbol to Type\n
\nallsubtypes(Integer)

2.4s

MLStyle (Julia)

12-element Array{Type,1}:\n Bool \n BigInt \n Int128 \n Int16 \n Int32 \n Int64 \n Int8 \n UInt128\n UInt16 \n UInt32 \n UInt64 \n UInt8

In order to enumerate every combination of concret types we can use Iterators.product :

collect(Iterators.product(
[Bool,BigInt]
,
[Bool,BigInt]
)
)

1.7s

MLStyle (Julia)

2×2 Array{Tuple{DataType,DataType},2}:\n (Bool, Bool) (Bool, BigInt) \n (BigInt, Bool) (BigInt, BigInt)

We also need a method that will take and abstract type annotation ( [ : T,:K] ) and replace it with the corresponding concrete subtypes:

function concretize!(args,
 concrete_args,
 idx)
\n
    for 
(j,K)
 in enumerate(concrete_args)
\n
        for i in idx[j]
\n
            args[i]
 = Symbol("$K")
\n
        end\n
    end\nend\n
\nargs = Symbol[:T,
 :K]
\nconcrete_args = 
(Bool,
 Bool)
\nidx = 
[
[1]
,
 
[2]
]
\n
\nconcretize!(args,
 concrete_args,
 idx)
\nargs

0.6s

MLStyle (Julia)

2-element Array{Symbol,1}:\n :Bool\n :Bool

Here idx contains the indices of where the abstract types annotations occurs in the method signature (they can occur at several positions, e.g. f(x::T, y::T, z::K) where {T,K} then idx = [[1,2],[3]] ).

We need to perform the same operation on our predicate (e.g. I -> I == T ), which can be an arbitrary expression. We will thus write a generic method that can replace a symbol by another in an expression:

#find all pair.first in expression and replace them by pair.second\nfunction ast_replace!(ex::Expr,
 pair::Pair)
\n
    for i=1:length(ex.args)
\n
        if ex.args[i]
 == pair.first\n
            ex.args[i]
 = pair.second\n
        else  
\n
            ast_replace!(ex.args[i]
,
 pair)
 
\n
        end\n
    end\n
    ex\nend 
\nast_replace!(ex::Symbol,
 pair)
 = nothing 
\nast_replace!(ex::LineNumberNode,
 pair)
 = nothing\n
\nast_replace!(:(I -> I == T)
,
 :T => Int8)

0.6s

MLStyle (Julia)

:(I->begin\n #= cell:15 =#\n I == Int8\n end)

We can now put all this together in a method that will generate the expression that performs the check for every combination of concrete parameters:

2.5s

MLStyle (Julia)

function check_types(fun,
 predicate,
 args,
 covs)
\n
  \t#setup\n
    out = Expr[
]
\n
    covs = match_subtype.(covs)
\n
    covs_symbols = 
[c.sub for c in covs]
\n
    idx  = 
[findall(args .== c)
 for c in covs_symbols]
\n
    
\n
    for p in Iterators.product(
[allsubtypes(c.sup)
 for c in covs]...)
 
\n
    
\n
    \t\t#build concrete arguments\n
        concrete_args = copy(args)
                
\n
\t\t\t  concretize!(concrete_args,
 p,
 idx)
\n
       
\n
        #build concrete predicate\n
        concrete_predicate = copy(predicate)
\n
        for pairs in 
[Pair(c,Symbol(t)
)
 for 
(c,t)
 in zip(covs_symbols,p)
]
\n
            ast_replace!(concrete_predicate,pairs)
\n
        end\n
        
\n
    \t\t#build final expression\n
        t = :(
 Tuple{$(concrete_args...)
}
 
)
\n
        I = :(
 inferred($fun,$t)
 
)
\n
        push!(out,
 quote 
\n
            report_type_check($concrete_predicate,$I,$fun,$concrete_args)
\n
        end)
\n
    end\n
    out\nend\n
\nfun,
 args = :add,
 
[:T,:K]
\ncovs = Any[:(T <: Integer)
,
 :(K <: Integer)
]
\nout_type = :(O -> O == T)
\n
\ncheck_types(fun,
 out_type,
 args,
 covs)
[1:2]

2.5s

MLStyle (Julia)

2-element Array{Expr,1}:\n quote\n #= cell:24 =#\n report_type_check((O->begin\n #= cell:32 =#\n O == Bool\n end), inferred(add, Tuple{Bool, Bool}), add, Symbol[:Bool, :Bool])\nend \n quote\n #= cell:24 =#\n report_type_check((O->begin\n #= cell:32 =#\n O == BigInt\n end), inferred(add, Tuple{BigInt, Bool}), add, Symbol[:BigInt, :Bool])\nend

In the generated code we call a helper function that will display the result of the type check instead of throwing an error:

function report_type_check(predicate,inferred,fun,types)
\n
    passed = predicate(inferred)
\n
    sig = join(types,",")
\n
    status = passed ? "SUCCESS" : "FAILURE"\n
    println(
 "$status: $(fun)($sig) \\t → $inferred" 
)
 
\nend

0.9s

MLStyle (Julia)

report_type_check (generic function with 1 method)

And finally we can update our checked macro:

macro checked(out,f)
\n
    sig = match_signature(quote $f end)
\n
    args = match_type_annotation.(sig.args)
\n
        
\n
    checks = check_types(Symbol(sig.name)
,out,args,sig.cov)
\n
    
\n
    esc(quote\n
        $f\n
        $(checks...)
\n
    end)
\nend\n
\n@checked O -> O == promote_type(T,K)
 function add(x::T,
 y::K)
 where 
{T<: Integer,
 K <: Integer}
 
\n
    x + y\nend

3.2s

MLStyle (Julia)

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n

All tests succeeded but the first one ( add(Bool,Bool) → Int64) since adding two booleans returns an Int64 . This could be fixed by modifying a bit the predicate to handle this particular case.

Limitations

One limitation of the implementation above is that it doesn't handle parametric types (e.g. Vector{T} where T <: Number ), but it should be relatively easy to generalize it to handle such a case. More general ones however might get tricky (e.g. Array{T,N} where {T <: Union{Number, String}, N} ).

The sheer number type combinations to check can also become an issue, for example there's 16 concrete subtype of Number in Base, so a function taking two numbers will generate 256 checks. There's 8998 concrete types in Base, so testing a function that has a single parameter annotated as Any would take that many checks, which is not very realistic.

For that reason this kind of tool might be better tailored for automatic tests, ran only during testing for some specific cases.

Final code

using MLStyle\nimport InteractiveUtils.subtypes\n
\nmatch_signature(ast)
 = @match ast begin\n
    quote\n
        $(::LineNumberNode)
\n
        function $name($(args...)
)
 where 
{$(cov...)
}
;
 $(block...)
 end\n
    end ||\n
    quote\n
        $(::LineNumberNode)
\n
        function $name($(args...)
)
;
 $(block...)
 end\n
    end && Do(cov = 
[
]
)
 =>\n
    
(name = name,
\n
     args = args,
\n
     cov  = cov)
\nend\n
\nmatch_type_annotation(ast)
 = @match ast begin :($(arg)::$(ann)
)
 => ann end\nmatch_subtype(ast)
 = @match ast begin :($(sub)<:$(sup)
)
 => 
(sub=sub,sup=sup)
 end\n
\nfunction inferred(f,args)
\n
    codeinfo = Base.code_typed(f,args)
[1]
\n
    codeinfo.second\nend\n
\nfunction allsubtypes(T::Type)
\n
    isconcretetype(T)
 && return 
[T]
\n
    t = Type[
]
\n
    map(K->allsubtypes(K,t)
,
 subtypes(T)
)
\n
    t\nend\nallsubtypes(T,t)
 = isconcretetype(T)
 ? push!(t,T)
 : map(K->allsubtypes(K,t)
,
 subtypes(T)
)
\nallsubtypes(T::Symbol)
 = allsubtypes(Core.eval(Main,T)
)#convert Symbol to Type\n
\nfunction concretize!(args,
 concrete_args,
 idx)
\n
    for 
(j,K)
 in enumerate(concrete_args)
\n
        for i in idx[j]
\n
            args[i]
 = Symbol("$K")
\n
        end\n
    end\nend\n
\nfunction ast_replace!(ex::Expr,
 pair::Pair)
\n
    for i=1:length(ex.args)
\n
        if ex.args[i]
 == pair.first\n
            ex.args[i]
 = pair.second\n
        else  
\n
            ast_replace!(ex.args[i]
,
 pair)
 
\n
        end\n
    end\n
    ex\nend 
\nast_replace!(ex::Symbol,
 pair)
 = nothing 
\nast_replace!(ex::LineNumberNode,
 pair)
 = nothing\n
\nfunction check_types(fun,
 predicate,
 args,
 covs)
\n
  \t#setup\n
    out = Expr[
]
\n
    covs = match_subtype.(covs)
\n
    covs_symbols = 
[c.sub for c in covs]
\n
    idx  = 
[findall(args .== c)
 for c in covs_symbols]
\n
    
\n
    for p in Iterators.product(
[allsubtypes(c.sup)
 for c in covs]...)
 
\n
    
\n
    \t\t#build concrete arguments\n
        concrete_args = copy(args)
                
\n
\t\t\t  concretize!(concrete_args,
 p,
 idx)
\n
       
\n
        #build concrete predicate\n
        concrete_predicate = copy(predicate)
\n
        for pairs in 
[Pair(c,Symbol(t)
)
 for 
(c,t)
 in zip(covs_symbols,p)
]
\n
            ast_replace!(concrete_predicate,pairs)
\n
        end\n
        
\n
    \t\t#build final expression\n
        t = :(
 Tuple{$(concrete_args...)
}
 
)
\n
        I = :(
 inferred($fun,$t)
 
)
\n
        push!(out,
 quote 
\n
            report_type_check($concrete_predicate,$I,$fun,$concrete_args)
\n
        end)
\n
    end\n
    out\nend\n
\nfunction report_type_check(predicate,inferred,fun,types)
\n
    passed = predicate(inferred)
\n
    sig = join(types,",")
\n
    status = passed ? "SUCCESS" : "FAILURE"\n
    println(
 "$status: $(fun)($sig) \\t → $inferred" 
)
 
\nend\n
\nmacro checked(out,f)
\n
    sig = match_signature(quote $f end)
\n
    args = match_type_annotation.(sig.args)
\n
        
\n
    checks = check_types(Symbol(sig.name)
,out,args,sig.cov)
\n
    
\n
    esc(quote\n
        $f\n
        $(checks...)
\n
    end)
\nend\n
\n@checked O -> O == promote_type(T,K)
 function add(x::T,
 y::K)
 where 
{T<: Integer,
 K <: Integer}
 
\n
    x + y\nend

Julia

"]]

Recommend

About Joyk

Aggregate valuable and interesting links.
Joyk means Joy of geeK

Adding static type checking to Julia in 100 lines of code

Recommend

在 IBM Cloud 上构建云就绪的 Express.js 应用程序

学习 Node.js，第 9 单元：单元测试

Google 研究员披露 Windows 10 0day 漏洞

研究称过去近三百年有 571 种植物灭绝

Chrome-derived browsers threaten to fork from Google, refuse to eliminate ad-blo...

日本网红眼药水被禁，可是我已经用了，还有救么？

大科学装置“运力”不足，用户们很捉急

43岁的他却有80岁的身体，这个基因给“衰老时钟”按了快进

如何写出无法维护的代码

asp.net core 集成JWT - 7tiny

About Joyk