Add 'StatsLearnModel'

src/StatsLearnModels.jl

@@ -8,7 +8,7 @@ using Tables
 using Distances
 using DataScienceTraits
 using StatsBase: mode, mean
-using ColumnSelectors: selector
+using ColumnSelectors: ColumnSelector, selector
 using TableTransforms: StatelessFeatureTransform
 
 import DataScienceTraits as DST

src/interface.jl

@@ -31,3 +31,22 @@ struct FittedModel{M,C}
 end
 
 Base.show(io::IO, ::FittedModel{M}) where {M} = print(io, "FittedModel{$(nameof(M))}")
+
+"""
+    StatsLearnModels.StatsLearnModel(model, incols, outcols)
+
+Wrapper type for learning models used for dispatch purposes.
+"""
+struct StatsLearnModel{M,I<:ColumnSelector,O<:ColumnSelector}
+  model::M
+  input::I
+  output::O
+end
+
+StatsLearnModel(model, incols, outcols) = StatsLearnModel(model, selector(incols), selector(outcols))
+
+function Base.show(io::IO, model::StatsLearnModel{M}) where {M}
+  println(io, "StatsLearnModel{$(nameof(M))}")
+  println(io, "├─ input: $(model.input)")
+  print(io, "└─ output: $(model.output)")
+end

test/runtests.jl

@@ -23,6 +23,14 @@ const SLM = StatsLearnModels
     @test sprint(show, fmodel) == "FittedModel{DecisionTreeClassifier}"
   end
 
+  @testset "StatsLearnModel" begin
+    model = SLM.StatsLearnModel(DecisionTreeClassifier(), [:a, :b], :c)
+    @test sprint(show, model) == """
+    StatsLearnModel{DecisionTreeClassifier}
+    ├─ input: [:a, :b]
+    └─ output: :c"""
+  end
+
   @testset "models" begin
     @testset "MLJ" begin
       Random.seed!(123)