From 7a93b093dda9c3ada33f318e744139df3ff1bab3 Mon Sep 17 00:00:00 2001
From: "Anthony D. Blaom" <anthony.blaom@gmail.com>
Date: Mon, 3 Jun 2024 16:13:06 +1200
Subject: [PATCH] add constructor trait; base docstring fallback on this if u
 can

---
 src/MLJModelInterface.jl |  1 +
 src/model_traits.jl      | 12 ++++++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/MLJModelInterface.jl b/src/MLJModelInterface.jl
index 6193110..244d141 100644
--- a/src/MLJModelInterface.jl
+++ b/src/MLJModelInterface.jl
@@ -33,6 +33,7 @@ const MODEL_TRAITS = [
     :reports_feature_importances,
     :deep_properties,
     :reporting_operations,
+    :constructor,
 ]
 
 const ABSTRACT_MODEL_SUBTYPES = [
diff --git a/src/model_traits.jl b/src/model_traits.jl
index 2269ecf..418cf30 100644
--- a/src/model_traits.jl
+++ b/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