Merge pull request #201 from JuliaAI/constructor

Add `constructor` trait

Author: ablaom

Project.toml

@@ -1,7 +1,7 @@
 name = "MLJModelInterface"
 uuid = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 authors = ["Thibaut Lienart and Anthony Blaom"]
-version = "1.9.6"
+version = "1.10.0"
 
 [deps]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -19,7 +19,7 @@ OrderedCollections = "1"
 Random = "<0.0.1, 1"
 ScientificTypes = "3"
 ScientificTypesBase = "3"
-StatisticalTraits = "3.2"
+StatisticalTraits = "3.3"
 Tables = "1"
 Test = "<0.0.1, 1"
 julia = "1.6"

src/MLJModelInterface.jl

@@ -33,6 +33,7 @@ const MODEL_TRAITS = [
     :reports_feature_importances,
     :deep_properties,
     :reporting_operations,
+    :constructor,
 ]
 
 const ABSTRACT_MODEL_SUBTYPES = [

src/model_traits.jl

@@ -13,10 +13,18 @@ const DeterministicDetector = Union{
 
 const StatTraits = StatisticalTraits
 
+# note that if F is a constructor, like `TunedModel`, then `docstring(F)` already falls
+# back to the function's document string.
 function StatTraits.docstring(M::Type{<:Model})
-    docstring = Base.Docs.doc(M) |> string
+    constructor = StatTraits.constructor(M)
+    # At time of writing, `constructor` is a new trait only overloaded for model wrappers
+    # that have multiple types associated with the same constructor (e.g., `TunedModel` is
+    # a constructor that can return objects of either `ProbabilisticTunedModel` or
+    # `DeterministicTunedModel` type. However, we want these bound to the same docstring.
+    C = isnothing(constructor) ? M : constructor
+    docstring = Base.Docs.doc(C) |> string
     if occursin("No documentation found", docstring)
-        docstring = synthesize_docstring(M)
+        docstring = synthesize_docstring(C)
     end
     return docstring
 end