Add a type parameter for AbstractStatsProcedure

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -6,6 +6,7 @@ using StatsBase
 @reexport using StatsModels
 
 import Base: ==, show
+import Base: eltype, getindex, iterate, length
 
 export @fieldequal,
        eachterm,

src/did.jl

@@ -1,16 +1,18 @@
 """
-    DiffinDiffsEstimator <: AbstractStatsProcedure
+    DiffinDiffsEstimator{T} <: AbstractStatsProcedure{T}
 
-Supertype for all types specifying the estimation procedure for difference-in-differences.
+Specify the estimation procedure for difference-in-differences.
 """
-abstract type DiffinDiffsEstimator <: AbstractStatsProcedure end
+struct DiffinDiffsEstimator{T} <: AbstractStatsProcedure{T} end
 
 """
     DefaultDID <: DiffinDiffsEstimator
 
 Default difference-in-differences estimator selected based on the context.
 """
-struct DefaultDID <: DiffinDiffsEstimator end
+const DefaultDID = DiffinDiffsEstimator{Tuple{}}
+
+show(io::IO, d::Type{DefaultDID}) = print(io, "DefaultDID")
 
 did(tr::AbstractTreatment, pr::AbstractParallel; kwargs...) =
     did(DefaultDID, tr, pr; kwargs...)

src/procedures.jl

@@ -1,35 +1,50 @@
 """
-    AbstractStatsProcedure
+    AbstractStatsProcedure{T<:NTuple{N,Function} where N}
 
 Supertype for all types specifying the procedure for statistical estimation or inference.
+
+The procedure is determined by the `parameters` of `T`,
+which are types of a sequence of functions.
 """
-abstract type AbstractStatsProcedure end
+abstract type AbstractStatsProcedure{T<:NTuple{N,Function} where N} end
+
+length(p::AbstractStatsProcedure{T}) where T = length(T.parameters)
+eltype(::Type{<:AbstractStatsProcedure}) = Function
+
+function getindex(p::AbstractStatsProcedure{T}, i) where T
+    fs = T.parameters[i]
+    return fs isa Type && fs <: Function ? fs.instance : [f.instance for f in fs]
+end
+
+iterate(p::AbstractStatsProcedure{T}, state=1) where T =
+    state > length(p) ? nothing : (p[state], state+1)
 
 """
-    StatsSpec{T<:AbstractStatsProcedure}
+    StatsSpec{T<:AbstractStatsProcedure, IsComplete}
 
-Record the specification for a statistical procedure and
-optionally a name for the specification.
+Record the specification for a statistical procedure of type `T`
+that may or may not be verified to be complete as indicated by `IsComplete`.
 
 The specification is recorded based on the arguments
 for a function that will conduct the procedure.
 It is assumed that a tuple of positional arguments accepted by the function
 can be constructed solely based on the type of each argument.
 
 # Fields
-- `name::String`: name for the specification.
+- `name::String`: an optional name for the specification.
 - `args::Dict{Symbol}`: positional arguments indexed based on their types.
 - `kwargs::Dict{Symbol}`: keyword arguments.
 """
-struct StatsSpec{T<:AbstractStatsProcedure}
+struct StatsSpec{T<:AbstractStatsProcedure, IsComplete}
     name::String
     args::Dict{Symbol}
     kwargs::Dict{Symbol}
-    StatsSpec(T::Type{<:AbstractStatsProcedure},
-        name::String, args::Dict{Symbol}, kwargs::Dict{Symbol}) =
-            new{T}(name, args, kwargs)
 end
 
+StatsSpec(T::Type{<:AbstractStatsProcedure}, name::String,
+    args::Dict{Symbol}, kwargs::Dict{Symbol}, IsComplete::Bool=false) =
+        StatsSpec{T,IsComplete}(name, args, kwargs)
+
 ==(a::StatsSpec{T}, b::StatsSpec{T}) where {T<:AbstractStatsProcedure} =
     a.args == b.args && a.kwargs == b.kwargs
 

test/did.jl

@@ -7,6 +7,24 @@ did(::Type{TestDID}, ::AbstractTreatment, ::AbstractParallel;
     yterm=:unknown, treatname=:unknown, treatintterms=nothing, xterms=nothing) =
         (yterm, treatname, treatintterms, xterms)
 
+@testset "DiffinDiffsEstimator" begin
+    d = DefaultDID()
+    @test length(d) == 0
+    @test eltype(d) == Function
+    @test_throws BoundsError d[1]
+    @test collect(d) == Function[]
+
+    d = TestDID()
+    @test length(d) == 2
+    @test eltype(d) == Function
+    @test d[1] == print
+    @test d[1:2] == [print, println]
+    @test d[[2,1]] == [println, print]
+    @test_throws BoundsError d[3]
+    @test_throws MethodError d[:a]
+    @test collect(d) == Function[print, println]
+end
+
 @testset "did wrapper" begin
     @testset "DefaultDID" begin
         @test_throws ErrorException did(TR, PR)

test/testutils.jl

@@ -1,5 +1,7 @@
 # Define simple generic types and methods for testing
 
+import Base: show
+
 struct TestTreatment <: AbstractTreatment
     time::Symbol
     ref::Int
@@ -14,8 +16,11 @@ end
 TestParallel(e::Int) = TestParallel{ParallelCondition,ParallelStrength}(e)
 tpara(c::ConstantTerm) = TestParallel{ParallelCondition,ParallelStrength}(c.n)
 
-struct NotImplemented <: DiffinDiffsEstimator end
-struct TestDID <: DiffinDiffsEstimator end
+const NotImplemented = DiffinDiffsEstimator{Tuple{typeof(println)}}
+show(io::IO, d::Type{NotImplemented}) = print(io, "NotImplemented")
+
+const TestDID = DiffinDiffsEstimator{Tuple{typeof(print), typeof(println)}}
+show(io::IO, d::Type{TestDID}) = print(io, "TestDID")
 
 const TR = TestTreatment(:t, 0)
 const PR = TestParallel(0)