From 922235ffc0a4b861456bc90bb14f054e4781c80b Mon Sep 17 00:00:00 2001
From: Phillip Alday <me@phillipalday.com>
Date: Sat, 2 Sep 2023 14:19:32 -0500
Subject: [PATCH] more docs more tests

---
 src/statsmodel.jl  | 11 +++++++++--
 test/statsmodel.jl | 20 ++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/statsmodel.jl b/src/statsmodel.jl
index 1c09ca5b..322bc972 100644
--- a/src/statsmodel.jl
+++ b/src/statsmodel.jl
@@ -150,8 +150,10 @@ termnames(t::FormulaTerm) = (termnames(t.lhs), termnames(t.rhs))
 """
     termnames(term::AbstractTerm)
 
-Return the name(s) of column(s) generated by a term.  Return value is either a
-`String` or an iterable of `String`s.
+Return the name(s) of column(s) generated by a term.  
+
+Return value is either a `String`, an iterable of `String`s or nothing if there
+no associated name (e.g. `termnames(InterceptTerm{false}())`).
 """
 termnames(::InterceptTerm{H}) where {H} = H ? "(Intercept)" : nothing
 termnames(t::ContinuousTerm) = string(t.sym)
@@ -160,6 +162,11 @@ termnames(t::Term) = string(t.sym)
 termnames(t::ConstantTerm) = string(t.n)
 termnames(t::FunctionTerm) = string(t.exorig)
 termnames(ts::TupleTerm) = mapreduce(termnames, vcat, ts)
+# these have some surprising behavior:
+# termnames(::InteractionTerm) always returns a vector
+# termnames(MatrixTerm(term(:a))) returns a scalar
+# termnames(MatrixTerm((term(a:), term(:b)))) returns a vector
+# but this is the same behavior as coefnames
 termnames(t::MatrixTerm) = mapreduce(termnames, vcat, t.terms)
 termnames(t::InteractionTerm) =
     kron_insideout((args...) -> join(args, " & "), vectorize.(termnames.(t.terms))...)
diff --git a/test/statsmodel.jl b/test/statsmodel.jl
index 639109af..c54183bc 100644
--- a/test/statsmodel.jl
+++ b/test/statsmodel.jl
@@ -234,14 +234,30 @@ Base.show(io::IO, m::DummyModTwo) = println(io, m.msg)
 
     m2 = fit(DummyModTwo, f, d)
     # make sure show() still works when there is no coeftable method
-    show(io, m2)
+    show(io, m2)          
+end
 
+@testset "termnames" begin
     # one final termnames check
     # note that `1` is still a ConstantTerm and not yet InterceptTerm
     # because apply_schema hasn't been called
     @test termnames(@formula(y ~ 1 + log(x) * y + (1+x|g)))[2] == 
           ["1", "log(x)", "y", "log(x) & y", "(1 + x) | g"]
-          
+    @test termnames(ConstantTerm(1)) == "1"
+    @test termnames(Term(:x)) == "x"
+    @test termnames(InterceptTerm{true}()) == "(Intercept)"
+    @test termnames(InterceptTerm{false}()) === nothing
+    @test termnames(ContinuousTerm(:x, 1, 0, 0, 0)) == "x"
+    cm = StatsModels.ContrastsMatrix([1 0; 0 1], ["b", "c"], ["a", "b", "c"], DummyCoding())
+    @test termnames(CategoricalTerm(:x, cm)) =="x"
+    @test termnames(FunctionTerm(log, [Term(:x)], :(log(x)))) == "log(x)"
+    # these next few seem a little weird but they're consistent with the 
+    # definition of coefnames
+    @test termnames(InteractionTerm(term.((:a, :b, :c)))) == ["a & b & c"]
+    @test termnames(MatrixTerm(term(:a))) == "a"
+    @test termnames(MatrixTerm((term(:a), term(:b)))) == ["a", "b"]
+    @test termnames((term(:a), term(:b))) == ["a", "b"]
+    @test termnames((term(:a),)) == "a"
 end
 
 @testset "lrtest" begin