Fix some issues of specset

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -7,9 +7,9 @@ using Reexport
 using StatsBase
 @reexport using StatsModels
 using StatsModels: TupleTerm
-using SplitApplyCombine: groupfind
+using SplitApplyCombine: groupfind, groupview
 using Tables: columntable, istable, rows, columns, getcolumn
-using TypedTables: Table, getproperty, getproperties
+using TypedTables: Table
 
 import Base: ==, show, union
 import Base: eltype, firstindex, lastindex, getindex, iterate, length

src/StatsProcedures.jl

@@ -322,6 +322,8 @@ while ignoring the orders.
 ≊(x::StatsSpec{A1,T}, y::StatsSpec{A2,T}) where {A1,A2,T} =
     x.args ≊ y.args
 
+_procedure(::StatsSpec{A,T}) where {A,T} = T
+
 function (sp::StatsSpec{A,T})(;
         verbose::Bool=false, keep=nothing, keepall::Bool=false) where {A,T}
     args = verbose ? merge(sp.args, (verbose=true,)) : sp.args
@@ -357,37 +359,57 @@ end
 
 function run_specset(sps::AbstractVector{<:StatsSpec};
         verbose::Bool=false, keep=nothing, keepall::Bool=false)
-    traces = Vector{NamedTuple}(undef, length(sps))
+    nsps = length(sps)
+    nsps == 0 && throw(ArgumentError("expect a nonempty vector"))
+    traces = Vector{NamedTuple}(undef, nsps)
     fill!(traces, NamedTuple())
     tb = Table(spec=sps, trace=traces)
-    gids = groupfind(r->procedure(r.spec), tb)
-    steps = pool((p() for p in keys(gids))...)
+    gids = groupfind(r->_procedure(r.spec)(), tb)
+    steps = pool((p for p in keys(gids))...)
+    ntask_total = 0
     for step in steps
-        verbose && printstyled("Running ", step, "\n", color=:green)
-        args = _specnames(step)
-        tras = _tracenames(step)
-        byf = r->merge(NamedTuple{args}(r.spec.args), NamedTuple{tras}(r.trace))
+        ntask = 0
+        verbose && printstyled("Running ", step, "...")
+        specn = _specnames(step)
+        tracn = _tracenames(step)
+        byf = r->merge(NamedTuple{specn}(r.spec.args), NamedTuple{tracn}(r.trace))
         taskids = vcat((gids[steps.procs[i]] for i in _sharedby(step))...)
         tasks = groupview(byf, view(tb, taskids))
         for (ins, subtb) in pairs(tasks)
             ret = _f(step)(ins...)
-            for tr in subtb.trace
-                tr = merge(tr, deepcopy(ret))
+            ntask += 1
+            ntask_total += 1
+            if ret !== nothing
+                for i in eachindex(subtb.trace)
+                    subtb.trace[i] = merge(subtb.trace[i], deepcopy(ret))
+                end
             end
         end
+        nprocs = length(_sharedby(step))
+        verbose && printstyled("Finished ", ntask, ntask > 1 ? " tasks" : " task", " for ",
+            nprocs, nprocs > 1 ? " procedures\n" : " procedure\n")
     end
+    nprocs = length(steps.procs)
+    verbose && printstyled("All steps finished (", ntask_total,
+        ntask_total > 1 ? " tasks" : " task", " for ", nprocs,
+        nprocs > 1 ? " procedures)\n" : " procedure)\n", bold=true, color=:green)
     if keepall
         return tb
     elseif keep===nothing
         return [r[end] for r in tb.trace]
     else
-        return [(;result=r.trace[end], (kv for kv in pairs(r.trace) if kv[1] in keep)...,
-            (kv for kv in pairs(r.spec.args) if kv[1] in keep)...) for r in tb]
+        if keep isa Symbol
+            keep = [keep]
+        elseif eltype(keep) != Symbol
+            throw(ArgumentError("expect Symbol or collections of Symbols for the value of option `keep`"))
+        end
+        return [(; (kv for kv in pairs(r.spec.args) if kv[1] in keep)...,
+            (kv for (i,kv) in enumerate(pairs(r.trace)) if
+            kv[1] in keep || i==length(r.trace))...,) for r in tb]
     end
 end
 
-function parse_specset_options(args)
-    options = :(Dict{Symbol, Any}())
+function _parse_kwargs!(options::Expr, args)
     for arg in args
         # Assume a symbol means the kwarg takes value true
         if isa(arg, Symbol)
@@ -397,19 +419,19 @@ function parse_specset_options(args)
             key = Expr(:quote, arg.args[1])
             push!(options.args, Expr(:call, :(=>), key, arg.args[2]))
         else
-            throw(ArgumentError("unexpected argument $arg to @specset"))
+            throw(ArgumentError("unexpected argument $arg"))
         end
     end
-    return options
 end
 
-function spec_walker(x, parsers, formatters)
-    @capture(x, StatsSpec(formatter_(parser_(rawargs__)))(;o__)) || return x
+function _spec_walker(x, parsers, formatters, ntargs_set)
+    @capture(x, StatsSpec(formatter_(parser_(rawargs__))...)(;o__)) || return x
     push!(parsers, parser)
     push!(formatters, formatter)
+    push!(ntargs_set, esc(:($parser($(rawargs...)))))
     length(o) > 0 &&
         @warn "[options] specified for individual StatsSpec are ignored inside @specset"
-    return :(push!(ntargs_set, $parser($(rawargs...))))
+    return :()
 end
 
 """
@@ -460,39 +482,35 @@ macro specset(args...)
     nargs = length(args)
     nargs == 0 && throw(ArgumentError("no argument is found for @specset"))
 
+    options = :(Dict{Symbol, Any}())
     if nargs > 1
         if isexpr(args[1], :vect, :hcat, :vcat)
-            options = parse_specset_options(args[1].args)
-            nargs > 2 && (default_args = args[2:end-1])
+            _parse_kwargs!(options, args[1].args)
+            nargs > 2 && (default_args = _args_kwargs(args[2:end-1]))
         else
-            default_args = args[1:end-1]
+            default_args = _args_kwargs(args[1:end-1])
         end
     else
-        options = :((;))
         default_args = nothing
     end
-
     specs = args[end]
     isexpr(specs, :block, :for) ||
         throw(ArgumentError("last argument to @specset must be begin/end block or for loop"))
 
-    parsers = []
-    formatters = []
-    blk = postwalk(x->spec_walker(x, parsers, formatters), specs)
-    length(parsers)==1 && length(formatters)==1 ||
+    parsers, formatters, ntargs_set = [], [], []
+    postwalk(x->_spec_walker(x, parsers, formatters, ntargs_set), specs)
+    length(unique!(parsers))==1 && length(unique!(formatters))==1 ||
         throw(ArgumentError("exactly one parser and one formatter are allowed for the inner @specset"))
 
-    parser = parsers[1]
-    formatter = formatters[1]
-    defaults = default_args === nothing ? :((;)) : :($parser($(default_args...)))
-    
-    return quote
-        local default_args = $defaults
-        local ntargs_set = NamedTuple[]
-        $blk
-        local nsps = length(ntargs_set)
-        local sps_set = [StatsSpec($formatter(merge(default_args, ntargs_set[i])))
-            for i in 1:nsps]
-        run_specset(sps_set; $(options...))
+    parser, formatter = parsers[1], formatters[1]
+    if default_args === nothing
+        defaults = :(NamedTuple())
+    else
+        defaults = esc(:($parser($(default_args[1]...); $(default_args[2]...))))
+    end
+    sps = :([])
+    for ntargs in ntargs_set
+        push!(sps.args, :(StatsSpec($(esc(formatter))($(esc(merge))($defaults, $ntargs))...)))
     end
+    return :(run_specset($sps; $options...))
 end

src/did.jl

@@ -94,7 +94,7 @@ The order of arguments is irrelevant.
 - `args... kwargs...`: a list of arguments to be processed by [`parse_didargs`](@ref) and [`valid_didargs`](@ref).
 """
 macro didspec(exprs...)
-    args, kwargs = args_kwargs(exprs)
+    args, kwargs = _args_kwargs(exprs)
     return esc(:(didspec($(args...); $(kwargs...))))
 end
 
@@ -111,8 +111,6 @@ function did(args...; verbose::Bool=false, keep=nothing, keepall::Bool=false, kw
     return sp(verbose=verbose, keep=keep, keepall=keepall)
 end
 
-const parse_did_options = parse_specset_options
-
 """
     @did [option option=val ...] "name" args... kwargs...
 
@@ -130,16 +128,16 @@ The options available are the same as the keyword arguments available for
 """
 macro did(args...)
     nargs = length(args)
+    options = :(Dict{Symbol, Any}())
     if nargs > 0 && isexpr(args[1], :vect, :hcat, :vcat)
-        options = parse_did_options(args[1].args)
+        _parse_kwargs!(options, args[1].args)
         if nargs > 1
             didargs = args[2:end]
         end
     else
-        options = :(Dict{Symbol, Any}())
         didargs = args
     end
-    dargs, dkwargs = args_kwargs(didargs)
+    dargs, dkwargs = _args_kwargs(didargs)
     return esc(:(StatsSpec(valid_didargs(parse_didargs($(dargs...); $(dkwargs...)))...)(; $options...)))
 end
 

src/utils.jl

@@ -94,14 +94,14 @@ macro unpack(functionname)
 end
 
 """
-    args_kwargs(exprs)
+    _args_kwargs(exprs)
 
 Partition a collection of expressions into two arrays
 such that all expressions in the second array has `head` being `:(=)`.
 This function is useful for separating out expressions
 for positional arguments and those for keyword arguments.
 """
-function args_kwargs(exprs)
+function _args_kwargs(exprs)
     args = []
     kwargs = []
     for expr in exprs

test/StatsProcedures.jl

@@ -1,10 +1,10 @@
 using DiffinDiffsBase: _f, _specnames, _tracenames,
-    _sharedby, _show_args
+    _sharedby, _show_args, _args_kwargs, _parse_kwargs!, _spec_walker, run_specset
 
-testvoidstep(a::String, b::String) = nothing
-const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep), (:a, :b), ()}
+testvoidstep(a::String) = nothing
+const TestVoidStep = StatsStep{:TestVoidStep, typeof(testvoidstep), (:a,), ()}
 
-testregstep(a::String, b::String) = (b=a, c=a*b,)
+testregstep(a::String, b::String) = (c=a*b,)
 const TestRegStep = StatsStep{:TestRegStep, typeof(testregstep), (:a, :b), ()}
 
 testlaststep(a::String, c::String) = (result=a*c,)
@@ -18,11 +18,11 @@ const TestInvalidStep = StatsStep{:TestInvalidStep, typeof(testinvalidstep), (:a
         @test sprint(show, TestVoidStep()) == "TestVoidStep"
         @test sprint(show, MIME("text/plain"), TestVoidStep()) == """
         TestVoidStep (StatsStep that calls testvoidstep):
-          arguments from StatsSpec: (:a, :b)
+          arguments from StatsSpec: (:a,)
           arguments from trace: ()"""
 
         @test _f(TestVoidStep()) == testvoidstep
-        @test _specnames(TestVoidStep()) == (:a, :b)
+        @test _specnames(TestVoidStep()) == (:a,)
         @test _tracenames(TestVoidStep()) == ()
 
         @test TestVoidStep()((a="a", b="b")) == (a="a", b="b")
@@ -39,7 +39,7 @@ const TestInvalidStep = StatsStep{:TestInvalidStep, typeof(testinvalidstep), (:a
         @test _specnames(TestRegStep()) == (:a, :b)
         @test _tracenames(TestRegStep()) == ()
 
-        @test TestRegStep()((a="a", b="b")) == (a="a", b="a", c="ab")
+        @test TestRegStep()((a="a", b="b")) == (a="a", b="b", c="ab")
     end
 
     @testset "TestLastStep" begin
@@ -209,7 +209,7 @@ end
     @test s1(keep=:a) == (a="a", result="aab")
     @test s1(keep=(:a,:c)) == (a="a", c="ab", result="aab")
     @test_throws ArgumentError s1(keep=1)
-    @test s1(keepall=true) == (a="a", b="a", c="ab", result="aab")
+    @test s1(keepall=true) == (a="a", b="b", c="ab", result="aab")
 
     s6 = StatsSpec("", NP, NamedTuple())
     @test s6() === nothing
@@ -221,3 +221,84 @@ end
 
     @test _show_args(stdout, s1) === nothing
 end
+
+function testparser(args...; kwargs...)
+    pargs = Pair{Symbol,Any}[kwargs...]
+    for arg in args
+        if arg isa Type{<:AbstractStatsProcedure}
+            push!(pargs, :p=>arg)
+        end
+    end
+    return (; pargs...)
+end
+
+testformatter(nt::NamedTuple) = (haskey(nt, :name) ? nt.name : "", nt.p, (a=nt.a, b=nt.b))
+
+@testset "specset" begin
+    @testset "run_specset" begin
+        s1 = StatsSpec("s1", RP, (a="a",b="b"))
+        s2 = StatsSpec("s2", RP, (a="a",b="b"))
+        s3 = StatsSpec("s3", RP, (a="a",b="b1"))
+        s4 = StatsSpec("s4", UP, (a="a",b="b"))
+        s5 = StatsSpec("s5", IP, (a="a",b="b"))
+        
+        @test run_specset([s1]) == ["aab"]
+        @test run_specset([s1,s2], verbose=true) == ["aab", "aab"]
+        @test run_specset([s1,s3], verbose=true) == ["aab", "aab1"]
+        @test run_specset([s1,s4], verbose=true) == ["aab", "ab"]
+        @test run_specset([s1,s5], verbose=true) == ["aab", "aab"]
+        @test run_specset([s1,s4,s5], verbose=true) == ["aab", "ab", "aab"]
+        @test_throws ArgumentError run_specset(StatsSpec[])
+
+        @test run_specset([s1], keep=:a) == [(a="a",result="aab")]
+        @test run_specset([s1], keep=[:a,:b]) == [(a="a", b="b", result="aab")]
+        @test run_specset([s1], keep=(:d,)) == [(result="aab",)]
+        @test run_specset([s1,s4], keep=[:a, :result]) ==
+            [(a="a", result="aab"), (a="a", c="ab")]
+         
+    end
+
+    @testset "_parse_kwargs!" begin
+        options = :(Dict{Symbol, Any}())
+        _parse_kwargs!(options, [:(a), :(b=1)])
+        @test eval(options) == Dict{Symbol, Any}(:a => true, :b => 1)
+        @test_throws ArgumentError _parse_kwargs!(options, [1])
+    end
+
+    @testset "@specset" begin
+        r = @specset a="a0" begin
+            StatsSpec(testformatter(testparser(RP; a="a1", b="b"))...)(;) end
+        @test r == ["a1a1b"]
+
+        r = @specset a="a0" begin
+            StatsSpec(testformatter(testparser(RP; b="b"))...)(;) end
+        @test r == ["a0a0b"]
+
+        r = @specset a="a0" b="b0" begin
+            StatsSpec(testformatter(testparser(RP))...)(;)
+            StatsSpec(testformatter(testparser(RP; a="a1", b="b1"))...)(;)
+        end
+        @test r == ["a0a0b0", "a1a1b1"]
+
+        r = @specset [verbose] a="a0" b="b0" begin
+            StatsSpec(testformatter(testparser(RP))...)(;)
+            StatsSpec(testformatter(testparser(RP; a="a1", c="c"))...)(;)
+        end
+        @test r == ["a0a0b0", "a1a1b0"]
+
+        a = "a0"
+        r = @specset [verbose] a=a begin
+            StatsSpec(testformatter(testparser(RP; b="b"))...)(;) end
+        @test r == ["a0a0b"]
+
+        r = []
+        for i in 1:3
+            a = "a"*string(i)
+            push!(r, @specset [verbose] a=a begin
+                StatsSpec(testformatter(testparser(RP; b="b"))...)(;)
+                StatsSpec(testformatter(testparser(RP; b="b1"))...)(;)
+            end)
+        end
+        @test r == [["a1a1b", "a1a1b1"], ["a2a2b", "a2a2b1"], ["a3a3b", "a3a3b1"]]
+    end
+end