Improve StatsSpec and fix an issue with macro did (#10)

junyuan-chen · web-flow · commit b3c99b2ff9bb · 2021-02-19T17:41:29.000-08:00
diff --git a/src/StatsProcedures.jl b/src/StatsProcedures.jl
@@ -54,10 +54,12 @@ See also [`required`](@ref) and [`default`](@ref).
 """
 transformed(::StatsStep, ::NamedTuple) = ()
 
-_get(a::NTuple{N,Symbol}, @nospecialize(nt::NamedTuple)) where N =
-    map(s->getfield(nt, s), a)
-_get(a::NamedTuple{S1}, nt::NamedTuple{S2}) where {S1,S2} =
-    map(s->getfield(ifelse(sym_in(s, S2), nt, a), s), S1)
+_get(nt::NamedTuple, args::Tuple{Vararg{Symbol}}) =
+    map(s->getfield(nt, s), args)
+
+# This method is useful for obtaining default values
+_get(nt::NamedTuple, args::NamedTuple) =
+    map(s->getfield(ifelse(sym_in(s, keys(nt)), nt, args), s), keys(args))
 
 """
     groupargs(s::StatsStep, ntargs::NamedTuple)
@@ -75,7 +77,7 @@ Instead, one should define methods for
 See also [`combinedargs`](@ref).
 """
 groupargs(s::StatsStep, @nospecialize(ntargs::NamedTuple)) =
-    (_get(required(s), ntargs)..., _get(default(s), ntargs)...,
+    (_get(ntargs, required(s))..., _get(ntargs, default(s))...,
     transformed(s, ntargs)...)
 
 """
@@ -162,6 +164,15 @@ function show(io::IO, ::MIME"text/plain", p::AbstractStatsProcedure{A,T}) where
     end
 end
 
+"""
+    _get_default(p::AbstractStatsProcedure, ntargs::NamedTuple)
+
+Obtain default values for arguments of each [`StatsStep`](@ref) in procedure `p`
+and then merge any default value for arguments not found in `ntargs` into `ntargs`.
+"""
+_get_default(p::AbstractStatsProcedure, @nospecialize(ntargs::NamedTuple)) =
+    merge((default(s) for s in p)..., ntargs)
+
 """
     SharedStatsStep
 
@@ -336,19 +347,20 @@ function show(io::IO, ::MIME"text/plain", p::PooledStatsProcedure)
 end
 
 """
-    StatsSpec{Alias, T<:AbstractStatsProcedure}
+    StatsSpec{T<:AbstractStatsProcedure}
 
 Record the specification for a statistical procedure of type `T`.
 
 An instance of `StatsSpec` is callable and
 its fields provide all information necessary for conducting the procedure.
-An optional name for the specification can be attached as parameter `Alias`.
+An optional name for the specification can be specified.
 
 # Fields
-- `args::NamedTuple`: arguments for the [`StatsStep`](@ref)s in `T`.
+- `name::String`: a name for the specification (takes `""` if not specified).
+- `args::NamedTuple`: arguments for the [`StatsStep`](@ref)s in `T` (default values are merged into `args` if not found in `args`).
 
 # Methods
-    (sp::StatsSpec{A,T})(; verbose::Bool=false, keep=nothing, keepall::Bool=false)
+    (sp::StatsSpec{T})(; verbose::Bool=false, keep=nothing, keepall::Bool=false)
 
 Execute the procedure of type `T` with the arguments specified in `args`.
 By default, a dedicated result object for `T` is returned if it is available.
@@ -359,37 +371,38 @@ Otherwise, the last value returned by the last [`StatsStep`](@ref) is returned.
 - `keep=nothing`: names (of type `Symbol`) of additional objects to be returned.
 - `keepall::Bool=false`: return all objects returned by each step.
 """
-struct StatsSpec{Alias, T<:AbstractStatsProcedure}
+struct StatsSpec{T<:AbstractStatsProcedure}
+    name::String
     args::NamedTuple
     StatsSpec(name::Union{Symbol,String},
         T::Type{<:AbstractStatsProcedure}, @nospecialize(args::NamedTuple)) =
-            new{Symbol(name),T}(args)
+            new{T}(string(name), args)
 end
 
 """
-    ==(x::StatsSpec{A1,T}, y::StatsSpec{A2,T})
+    ==(x::StatsSpec{T}, y::StatsSpec{T})
 
 Test whether two instances of [`StatsSpec`](@ref)
 with the same parameter `T` also have the same field `args`.
 
 See also [`≊`](@ref).
 """
-==(x::StatsSpec{A1,T}, y::StatsSpec{A2,T}) where {A1,A2,T} = x.args == y.args
+==(x::StatsSpec{T}, y::StatsSpec{T}) where T = x.args == y.args
 
 """
-    ≊(x::StatsSpec{A1,T}, y::StatsSpec{A2,T})
+    ≊(x::StatsSpec{T}, y::StatsSpec{T})
 
 Test whether two instances of [`StatsSpec`](@ref)
 with the same parameter `T` also have the field `args`
 containing the same sets of key-value pairs
 while ignoring the orders.
 """
-≊(x::StatsSpec{A1,T}, y::StatsSpec{A2,T}) where {A1,A2,T} = x.args ≊ y.args
+≊(x::StatsSpec{T}, y::StatsSpec{T}) where T = x.args ≊ y.args
 
-_procedure(::StatsSpec{A,T}) where {A,T} = T
+_procedure(::StatsSpec{T}) where T = T
 
-function (sp::StatsSpec{A,T})(;
-        verbose::Bool=false, keep=nothing, keepall::Bool=false) where {A,T}
+function (sp::StatsSpec{T})(;
+        verbose::Bool=false, keep=nothing, keepall::Bool=false) where T
     args = verbose ? merge(sp.args, (verbose=true,)) : sp.args
     ntall = result(T, foldl(|>, T(), init=args))
     if keepall
@@ -416,12 +429,12 @@ function (sp::StatsSpec{A,T})(;
     end
 end
 
-show(io::IO, ::StatsSpec{A}) where {A} = print(io, A==Symbol("") ? "unnamed" : A)
+show(io::IO, sp::StatsSpec) = print(io, sp.name=="" ? "unnamed" : sp.name)
 
 _show_args(::IO, ::StatsSpec) = nothing
 
-function show(io::IO, ::MIME"text/plain", sp::StatsSpec{A,T}) where {A,T}
-    print(io, A==Symbol("") ? "unnamed" : A, " (", typeof(sp).name.name,
+function show(io::IO, ::MIME"text/plain", sp::StatsSpec{T}) where T
+    print(io, sp.name=="" ? "unnamed" : sp.name, " (", typeof(sp).name.name,
         " for ", T.parameters[1], ")")
     _show_args(io, sp)
 end
@@ -517,19 +530,18 @@ function proceed(sps::AbstractVector{<:StatsSpec};
 end
 
 function _parse!(options::Expr, args)
-    noproceed = false
     for arg in args
         # Assume a symbol means the kwarg takes value true
         if isa(arg, Symbol)
             if arg == :noproceed
-                noproceed = true
+                return true
             else
                 key = Expr(:quote, arg)
                 push!(options.args, Expr(:call, :(=>), key, true))
             end
         elseif isexpr(arg, :(=))
             if arg.args[1] == :noproceed
-                noproceed = arg.args[2]
+                return arg.args[2]
             else
                 key = Expr(:quote, arg.args[1])
                 push!(options.args, Expr(:call, :(=>), key, arg.args[2]))
@@ -538,18 +550,24 @@ function _parse!(options::Expr, args)
             throw(ArgumentError("unexpected option $arg"))
         end
     end
-    return noproceed
+    return false
 end
 
 function _spec_walker1(x, parsers, formatters, ntargs_set)
-    @capture(x, StatsSpec(formatter_(parser_(rawargs__))...)(;o__)) || return x
+    @capture(x, spec_(formatter_(parser_(rawargs__))...)(;o__)) || return x
+    # spec may be a GlobalRef
+    name = spec isa Symbol ? spec : spec.name
+    name == :StatsSpec || return x
     push!(parsers, parser)
     push!(formatters, formatter)
     return :(push!($ntargs_set, $parser($(rawargs...))))
 end
 
 function _spec_walker2(x, parsers, formatters, ntargs_set)
-    @capture(x, StatsSpec(formatter_(parser_(rawargs__))...)) || return x
+    @capture(x, spec_(formatter_(parser_(rawargs__))...)) || return x
+    # spec may be a GlobalRef
+    name = spec isa Symbol ? spec : spec.name
+    name == :StatsSpec || return x
     push!(parsers, parser)
     push!(formatters, formatter)
     return :(push!($ntargs_set, $parser($(rawargs...))))
@@ -613,7 +631,8 @@ macro specset(args...)
     isexpr(specs, :block, :for) ||
         throw(ArgumentError("last argument to @specset must be begin/end block or for loop"))
 
-    parsers, formatters, ntargs_set = Symbol[], Symbol[], NamedTuple[]
+    # parser and formatter may be GlobalRef
+    parsers, formatters, ntargs_set = [], [], NamedTuple[]
     walked = postwalk(x->_spec_walker1(x, parsers, formatters, ntargs_set), specs)
     walked = postwalk(x->_spec_walker2(x, parsers, formatters, ntargs_set), walked)
     nparser = length(unique!(parsers))
diff --git a/src/did.jl b/src/did.jl
@@ -87,7 +87,7 @@ with the specified arguments.
 """
 didspec(args...; kwargs...) = StatsSpec(_didargs(args...; kwargs...)...)
 
-function _show_args(io::IO, sp::StatsSpec{A,<:DiffinDiffsEstimator}) where A
+function _show_args(io::IO, sp::StatsSpec{<:DiffinDiffsEstimator})
     if haskey(sp.args, :tr) || haskey(sp.args, :pr)
         print(io, ":")
         haskey(sp.args, :tr) && print(io, "\n  ", sp.args[:tr])
@@ -139,7 +139,7 @@ macro did(args...)
     nargs = length(args)
     options = :(Dict{Symbol, Any}())
     noproceed = false
-    didargs = []
+    didargs = ()
     if nargs > 0
         if isexpr(args[1], :vect, :hcat, :vcat)
             noproceed = _parse!(options, args[1].args)
@@ -150,9 +150,9 @@ macro did(args...)
     end
     dargs, dkwargs = _args_kwargs(didargs)
     if noproceed
-        return esc(:(StatsSpec(valid_didargs(parse_didargs($(dargs...); $(dkwargs...)))...)))
+        return :(StatsSpec(valid_didargs(parse_didargs($(esc.(dargs)...); $(esc.(dkwargs)...)))...))
     else
-        return esc(:(StatsSpec(valid_didargs(parse_didargs($(dargs...); $(dkwargs...)))...)(; $options...)))
+        return :(StatsSpec(valid_didargs(parse_didargs($(esc.(dargs)...); $(esc.(dkwargs)...)))...)(; $(esc(options))...))
     end
 end
 
diff --git a/test/StatsProcedures.jl b/test/StatsProcedures.jl
@@ -1,4 +1,4 @@
-using DiffinDiffsBase: _f, _get, groupargs,
+using DiffinDiffsBase: _f, _get, groupargs, _get_default,
     _sharedby, _show_args, _args_kwargs, _parse!, pool, proceed
 import DiffinDiffsBase: required, default, transformed, combinedargs
 
@@ -26,20 +26,20 @@ combinedargs(::TestCombineStep, ntargs) = [nt.b for nt in ntargs]
 
 testinvalidstep(a::String, b::String) = b, false
 const TestInvalidStep = StatsStep{:TestInvalidStep, typeof(testinvalidstep)}
-default(::TestInvalidStep) = (a="a",b="b")
+default(::TestInvalidStep) = (a="a", b="b")
 
 const TestUnnamedStep = StatsStep{:TestUnnamedStep, typeof(testinvalidstep)}
 
 @testset "StatsStep" begin
     @testset "_get" begin
-        @test _get((), NamedTuple()) == ()
-        @test _get((:a,), (a=1, b=2)) == (1,)
-        @test_throws ErrorException _get((:a,), (b=2,))
+        @test _get(NamedTuple(), ()) == ()
+        @test _get((a=1, b=2), (:a,)) == (1,)
+        @test_throws ErrorException _get((b=2,), (:a,))
 
         @test _get(NamedTuple(), NamedTuple()) == ()
-        @test _get((a=1,), (b=2,)) == (1,)
-        @test _get((a=1,), (a=1, b=2)) == (1,)
-        @test _get((a=1, b=2), (a=2,)) == (2, 2)
+        @test _get((a=1,), (b=2,)) == (2,)
+        @test _get((a=1,), (a=2, b=2)) == (1, 2)
+        @test _get((a=1, b=2), (a=2,)) == (1,)
     end
 
     @testset "args" begin
@@ -126,6 +126,12 @@ const ep = EP()
         NullProcedure (TestProcedure with 0 step)"""
 end
 
+@testset "_get_default" begin
+    @test _get_default(rp, NamedTuple()) == (a="a", b="b")
+    @test _get_default(rp, (a="a1",)) == (a="a1", b="b")
+    @test _get_default(rp, (c="c",)) == (a="a", b="b", c="c")
+end
+
 @testset "SharedStatsStep" begin
     s1 = SharedStatsStep(TestRegStep(), 1)
     s2 = SharedStatsStep(TestRegStep(), [3,2])
