From 3bd2966540b311c59a514836393fe90d59ac712e Mon Sep 17 00:00:00 2001
From: "Anthony D. Blaom" <anthony.blaom@gmail.com>
Date: Fri, 2 Aug 2024 12:04:27 +1200
Subject: [PATCH] fix is_wrapper bug to close #24

---
 Project.toml                  |  2 +-
 src/FeatureSelection.jl       |  1 +
 src/models/featureselector.jl | 10 ----------
 src/models/rfe.jl             | 13 -------------
 src/shared.jl                 | 16 ++++++++++++++++
 test/models/rfe.jl            |  8 ++++++++
 6 files changed, 26 insertions(+), 24 deletions(-)
 create mode 100644 src/shared.jl

diff --git a/Project.toml b/Project.toml
index d5e596b..619d2ae 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "FeatureSelection"
 uuid = "33837fe5-dbff-4c9e-8c2f-c5612fe2b8b6"
 authors = ["Anthony D. Blaom <anthony.blaom@gmail.com>", "Samuel Okon <okonsamuel50@gmail.com"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
diff --git a/src/FeatureSelection.jl b/src/FeatureSelection.jl
index e097c4d..a2e8a54 100644
--- a/src/FeatureSelection.jl
+++ b/src/FeatureSelection.jl
@@ -9,5 +9,6 @@ const MMI = MLJModelInterface
 ## Includes
 include("models/featureselector.jl")
 include("models/rfe.jl")
+include("shared.jl")
 
 end # module
diff --git a/src/models/featureselector.jl b/src/models/featureselector.jl
index f314b5f..470de9b 100644
--- a/src/models/featureselector.jl
+++ b/src/models/featureselector.jl
@@ -87,16 +87,6 @@ MMI.metadata_model(
     load_path = "FeatureSelection.FeatureSelector"
 )
 
-## Pkg Traits
-MMI.metadata_pkg(
-    FeatureSelector,
-    package_name = "FeatureSelection",
-    package_uuid = "33837fe5-dbff-4c9e-8c2f-c5612fe2b8b6",
-    package_url = "https://github.com/JuliaAI/FeatureSelection.jl",
-    is_pure_julia = true,
-    package_license = "MIT"
-)
-
 ## Docstring
 """
 $(MMI.doc_header(FeatureSelector))
diff --git a/src/models/rfe.jl b/src/models/rfe.jl
index 2c53b7a..10a49fe 100644
--- a/src/models/rfe.jl
+++ b/src/models/rfe.jl
@@ -453,16 +453,3 @@ end
 function MMI.training_losses(model::RFE, rfe_report)
     return MMI.training_losses(model.model, rfe_report.model_report)
 end
-
-## Pkg Traits
-MMI.metadata_pkg.(
-    (
-        DeterministicRecursiveFeatureElimination,
-        ProbabilisticRecursiveFeatureElimination,
-    ),
-    package_name       = "FeatureSelection",
-    package_uuid       = "33837fe5-dbff-4c9e-8c2f-c5612fe2b8b6",
-    package_url        = "https://github.com/JuliaAI/FeatureSelection.jl",
-    is_pure_julia      = true,
-    package_license    = "MIT"
-)
diff --git a/src/shared.jl b/src/shared.jl
new file mode 100644
index 0000000..31c6307
--- /dev/null
+++ b/src/shared.jl
@@ -0,0 +1,16 @@
+## Pkg Traits
+
+for M in [
+    :FeatureSelector,
+    :DeterministicRecursiveFeatureElimination,
+    :ProbabilisticRecursiveFeatureElimination,
+    ]
+    quote
+        MMI.package_name(::Type{<:$M})  = "FeatureSelection"
+        MMI.package_uuid(::Type{<:$M})  = "33837fe5-dbff-4c9e-8c2f-c5612fe2b8b6"
+        MMI.package_url(::Type{<:$M})   = "https://github.com/JuliaAI/FeatureSelection.jl"
+        MMI.is_pure_julia(::Type{<:$M}) = true
+        MMI.package_license(::Type{<:$M})   = "MIT"
+    end |> eval
+end
+
diff --git a/test/models/rfe.jl b/test/models/rfe.jl
index c257936..4c5e2d0 100644
--- a/test/models/rfe.jl
+++ b/test/models/rfe.jl
@@ -191,3 +191,11 @@ end
 
     @test predict(mach2, (; x1=rand(2), x2=rand(2))) == yhat
 end
+
+@testset "is_wrapper" begin
+    m = RecursiveFeatureElimination(DecisionTreeClassifier())
+    @test is_wrapper(m)
+    @test is_wrapper(typeof(m))
+end
+
+true