Improve methods of Base.show

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -1,6 +1,6 @@
 module DiffinDiffsBase
 
-using CSV
+using CSV: File
 using Reexport
 using StatsBase
 @reexport using StatsModels
@@ -14,6 +14,7 @@ export @fieldequal,
        kwarg,
        @unpack,
        exampledata,
+       sprintcompact,
 
        EleOrVec,
        TreatmentSharpness,
@@ -44,8 +45,10 @@ export @fieldequal,
        parse_treat,
 
        AbstractDiffinDiffs,
+       DefaultDID,
        did,
        DIDSpec,
+       isnamed,
        spec,
        @spec,
        @did,

src/did.jl

@@ -134,6 +134,20 @@ end
 ==(a::DIDSpec{T}, b::DIDSpec{T}) where {T<:AbstractDiffinDiffs} =
     a.args == b.args && a.kwargs == b.kwargs
 
+isnamed(sp::DIDSpec) = sp.name != ""
+
+function show(io::IO, sp::DIDSpec{T}) where {T}
+    print(io, "DIDSpec{", sprintcompact(T), "}")
+    isnamed(sp) && print(io, ": ", sp.name)
+    if get(io, :compact, false) || !haskey(sp.args, :tr) && !haskey(sp.args, :pr)
+        return
+    elseif !isnamed(sp)
+        print(io, ":")    
+    end
+    haskey(sp.args, :tr) && print(io, "\n  ", sprintcompact(sp.args[:tr]))
+    haskey(sp.args, :pr) && print(io, "\n  ", sprintcompact(sp.args[:pr]))
+end
+
 """
     spec(args...; kwargs...)
 

src/parallels.jl

@@ -111,8 +111,8 @@ function show(io::IO, pr::NeverTreatedParallel)
         print(io, "NeverTreated{", pr.c, ",", pr.s, "}(", pr.e,")")
     else
         println(io, pr.s, " trends with any never-treated group:")
-        println(io, "  Never-treated groups: ", pr.e)
-        pr.c isa Unconditional ? nothing : println(io, "  ", pr.c)
+        print(io, "  Never-treated groups: ", pr.e)
+        pr.c isa Unconditional || print(io, "\n  ", pr.c)
     end
 end
 
@@ -193,8 +193,8 @@ function show(io::IO, pr::NotYetTreatedParallel)
     else
         println(io, pr.s, " trends with any not-yet-treated group:")
         println(io, "  Not-yet-treated groups: ", pr.e)
-        println(io, "  Treated since: ", pr.emin isa Nothing ? "not specified" : pr.emin)
-        pr.c isa Unconditional ? nothing : println(io, "  ", pr.c)
+        print(io, "  Treated since: ", pr.emin isa Nothing ? "not specified" : pr.emin)
+        pr.c isa Unconditional || print(io, "\n  ", pr.c)
     end
 end
 

src/treatments.jl

@@ -72,7 +72,7 @@ function show(io::IO, tr::DynamicTreatment)
     else
         println(io, tr.s, " dynamic treatment:")
         println(io, "  column name of time variable: ", tr.time)
-        println(io, "  excluded relative time: ", tr.exc)
+        print(io, "  excluded relative time: ", tr.exc)
     end
 end
 

src/utils.jl

@@ -110,5 +110,7 @@ Return a `CSV.File` by loading the example dataset with the specified name.
 function exampledata(name::Union{Symbol,String})
     "$(name)" in exampledata() ||
         throw(ArgumentError("example dataset $(name) is not found"))
-    return CSV.File((@__DIR__)*"/../data/$(name).csv")
+    return File((@__DIR__)*"/../data/$(name).csv")
 end
+
+sprintcompact(x) = sprint(show, x; context=:compact=>true)

test/did.jl

@@ -1,12 +1,12 @@
 using DiffinDiffsBase: parse_didargs
 import DiffinDiffsBase: did
 
+const testterm = treat(:g, TR, PR)
+
 did(::Type{TestDID}, ::AbstractTreatment, ::AbstractParallel;
     yterm=:unknown, treatname=:unknown, treatintterms=nothing, xterms=nothing) =
         (yterm, treatname, treatintterms, xterms)
 
-const testterm = treat(:g, TR, PR)
-
 @testset "did wrapper" begin
     @testset "DefaultDID" begin
         @test_throws ErrorException did(TR, PR)
@@ -86,6 +86,45 @@ end
     @test_throws ArgumentError parse_didargs('a', :a, 1)
 end
 
+@testset "DIDSpec" begin
+    @testset "show" begin
+        sp = DIDSpec("", Dict{Symbol,Any}(), Dict{Symbol,Any}())
+        @test sprint(show, sp) == "DIDSpec{DefaultDID}"
+        @test sprintcompact(sp) == "DIDSpec{DefaultDID}"
+
+        sp = DIDSpec("name", Dict{Symbol,Any}(), Dict{Symbol,Any}())
+        @test sprint(show, sp) == "DIDSpec{DefaultDID}: name"
+        @test sprintcompact(sp) == "DIDSpec{DefaultDID}: name"
+        
+        sp = DIDSpec("name", Dict(:d=>TestDID, :tr=>TR, :pr=>PR), Dict(:a=>1, :b=>2))
+        @test sprint(show, sp) == """
+            DIDSpec{TestDID}: name
+              TestTreatment(:t, 0)
+              TestParallel{ParallelCondition,ParallelStrength}(0)"""
+        @test sprintcompact(sp) == "DIDSpec{TestDID}: name"
+        
+        sp = DIDSpec("", Dict(:d=>TestDID, :tr=>dynamic(:time,-1), :pr=>nevertreated(-1)),
+            Dict{Symbol,Any}())
+        @test sprint(show, sp) == """
+            DIDSpec{TestDID}:
+              Dynamic{S}(-1)
+              NeverTreated{U,P}([-1])"""
+        @test sprintcompact(sp) == "DIDSpec{TestDID}"
+        
+        sp = DIDSpec("", Dict(:d=>TestDID, :tr=>dynamic(:time,-1)), Dict{Symbol,Any}())
+        @test sprint(show, sp) == """
+            DIDSpec{TestDID}:
+              Dynamic{S}(-1)"""
+        @test sprintcompact(sp) == "DIDSpec{TestDID}"
+        
+        sp = DIDSpec("", Dict(:d=>TestDID, :pr=>nevertreated(-1)), Dict{Symbol,Any}())
+        @test sprint(show, sp) == """
+            DIDSpec{TestDID}:
+              NeverTreated{U,P}([-1])"""
+        @test sprintcompact(sp) == "DIDSpec{TestDID}"
+    end
+end
+
 @testset "spec @spec" begin
     testname = "name"
 

test/parallels.jl

@@ -6,10 +6,10 @@
 
     @testset "show" begin
         @test sprint(show, Unconditional()) == "Unconditional"
-        @test sprint(show, Unconditional(); context=:compact => true) == "U"
+        @test sprintcompact(Unconditional()) == "U"
 
         @test sprint(show, Exact()) == "Parallel"
-        @test sprint(show, Exact(); context=:compact => true) == "P"
+        @test sprintcompact(Exact()) == "P"
     end
 end
 
@@ -68,17 +68,13 @@ end
     @testset "show" begin
         @test sprint(show, nt0) == """
             Parallel trends with any never-treated group:
-              Never-treated groups: [0]
-            """
-        @test sprint(show, nt0; context=:compact => true) ==
-            "NeverTreated{U,P}([0])"
+              Never-treated groups: [0]"""
+        @test sprintcompact(nt0) == "NeverTreated{U,P}([0])"
 
         @test sprint(show, nt1) == """
             Parallel trends with any never-treated group:
-              Never-treated groups: [0, 1]
-            """
-        @test sprint(show, nt1; context=:compact => true) ==
-            "NeverTreated{U,P}([0, 1])"
+              Never-treated groups: [0, 1]"""
+        @test sprintcompact(nt1) == "NeverTreated{U,P}([0, 1])"
     end
 end
 
@@ -152,33 +148,25 @@ end
         @test sprint(show, ny0) == """
             Parallel trends with any not-yet-treated group:
               Not-yet-treated groups: [0]
-              Treated since: not specified
-            """
-        @test sprint(show, ny0; context=:compact => true) ==
-            "NotYetTreated{U,P}([0], NA)"
+              Treated since: not specified"""
+        @test sprintcompact(ny0) == "NotYetTreated{U,P}([0], NA)"
 
         @test sprint(show, ny1) == """
             Parallel trends with any not-yet-treated group:
               Not-yet-treated groups: [0, 1]
-              Treated since: not specified
-            """
-        @test sprint(show, ny1; context=:compact => true) ==
-            "NotYetTreated{U,P}([0, 1], NA)"
+              Treated since: not specified"""
+        @test sprintcompact(ny1) == "NotYetTreated{U,P}([0, 1], NA)"
 
         @test sprint(show, ny2) == """
             Parallel trends with any not-yet-treated group:
               Not-yet-treated groups: [0, 1]
-              Treated since: [0]
-            """
-        @test sprint(show, ny2; context=:compact => true) ==
-            "NotYetTreated{U,P}([0, 1], [0])"
+              Treated since: [0]"""
+        @test sprintcompact(ny2) == "NotYetTreated{U,P}([0, 1], [0])"
 
         @test sprint(show, ny3) == """
             Parallel trends with any not-yet-treated group:
               Not-yet-treated groups: [0, 1]
-              Treated since: [0, 1]
-            """
-        @test sprint(show, ny3; context=:compact => true) ==
-            "NotYetTreated{U,P}([0, 1], [0, 1])"
+              Treated since: [0, 1]"""
+        @test sprintcompact(ny3) == "NotYetTreated{U,P}([0, 1], [0, 1])"
     end
 end

test/treatments.jl

@@ -5,7 +5,7 @@
 
     @testset "show" begin
         @test sprint(show, SharpDesign()) == "Sharp"
-        @test sprint(show, SharpDesign(); context=:compact => true) == "S"
+        @test sprintcompact(SharpDesign()) == "S"
     end
 end
 
@@ -46,16 +46,12 @@ end
         @test sprint(show, dt0) == """
             Sharp dynamic treatment:
               column name of time variable: month
-              excluded relative time: -1
-            """
-        @test sprint(show, dt0; context=:compact => true) ==
-            "Dynamic{S}(-1)"
+              excluded relative time: -1"""
+        @test sprintcompact(dt0) == "Dynamic{S}(-1)"
         @test sprint(show, dt1) == """
             Sharp dynamic treatment:
               column name of time variable: month
-              excluded relative time: [-2, -1]
-            """
-        @test sprint(show, dt1; context=:compact => true) ==
-            "Dynamic{S}([-2, -1])"
+              excluded relative time: [-2, -1]"""
+        @test sprintcompact(dt1) == "Dynamic{S}([-2, -1])"
     end
 end