From a00c55fb3fc85ee5105af49eddebcd42f3250fd8 Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Tue, 15 Mar 2022 15:04:54 +0100
Subject: [PATCH] Support attaching documentation to the enum-module, and
 individual instances.

---
 README.md        | 28 ++++++++++++++++++++-----
 src/EnumX.jl     | 22 +++++++++++++++++++-
 test/runtests.jl | 53 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 0ec21bf..4ad7140 100644
--- a/README.md
+++ b/README.md
@@ -56,13 +56,30 @@ Since the only reserved name in the example above is the module `Fruit` we can c
 another enum with overlapping instance names (this would not be possible with `Base.@enum`):
 
 ```julia
-julia> @enumx YellowFruits Banana Lemon
+julia> @enumx YellowFruit Banana Lemon
 
-julia> YellowFruits.Banana
-YellowFruits.Banana = 0
+julia> YellowFruit.Banana
+YellowFruit.Banana = 0
 ```
 
-`@enumx` also allows for duplicate values:
+Instances can be documented like `struct` fields. A docstring before the macro is
+attached to the *module* `Fruit` (i.e. not the "hidden" type `Fruit.T`):
+
+```julia
+julia> "Documentation for Fruit enum-module."
+       @enumx Fruit begin
+           "Documentation for Fruit.Apple instance."
+           Apple
+       end
+
+help?> Fruit
+  Documentation for Fruit enum-module.
+
+help?> Fruit.Apple
+  Documentation for Fruit.Apple instance.
+```
+
+`@enumx` allows for duplicate values (unlike `Base.@enum`):
 
 ```julia
 julia> @enumx Fruit Apple=1 Banana=1
@@ -79,7 +96,8 @@ julia> Fruit.Banana
 Fruit.Apple = Fruit.Banana = 1
 ```
 
-`@enumx` also lets you use previous enum names for value initialization:
+`@enumx` lets you use previous enum names for value initialization:
+
 ```julia
 julia> @enumx Fruit Apple Banana Orange=Apple
 
diff --git a/src/EnumX.jl b/src/EnumX.jl
index 1d73b4b..7e7945f 100644
--- a/src/EnumX.jl
+++ b/src/EnumX.jl
@@ -39,10 +39,18 @@ function enumx(_module_, args)
         syms = args
     end
     name_value_map = Vector{Pair{Symbol, baseT}}()
+    doc_entries = Vector{Pair{Symbol,Expr}}()
     next = zero(baseT)
     first = true
     for s in syms
         s isa LineNumberNode && continue
+        # Handle doc expressions
+        doc_expr = nothing
+        if Meta.isexpr(s, :macrocall, 4) && s.args[1] isa GlobalRef && s.args[1].mod === Core &&
+               s.args[1].name === Symbol("@doc")
+            doc_expr = s
+            s = s.args[4]
+        end
         if s isa Symbol
             if !first && next == typemin(baseT)
                 panic("value overflow for Enum $(modname): $(modname).$(s) = $(next).")
@@ -78,6 +86,11 @@ function enumx(_module_, args)
             )
         end
         push!(name_value_map, sym => next)
+        if doc_expr !== nothing
+            # Replace the documented expression since it might be :(Apple = ...)
+            doc_expr.args[4] = sym
+            push!(doc_entries, sym => doc_expr)
+        end
 
         next += oneunit(baseT)
         first = false
@@ -103,7 +116,14 @@ function enumx(_module_, args)
             Expr(:const, Expr(:(=), esc(k), Expr(:call, esc(T), v)))
         )
     end
-    return Expr(:toplevel, Expr(:module, false, esc(modname), module_block), nothing)
+    for (_, v) in doc_entries
+        push!(module_block.args, v)
+    end
+    # Document the module and the type
+    mdoc = Expr(:block, Expr(:meta, :doc), esc(modname))
+    # TODO: Attach to the type too?
+    # Tdoc = Expr(:block, Expr(:meta, :doc), Expr(:., esc(modname), QuoteNode(T)))
+    return Expr(:toplevel, Expr(:module, false, esc(modname), module_block), mdoc, #=Tdoc,=# nothing)
 end
 
 function Base.show(io::IO, ::MIME"text/plain", x::E) where E <: Enum
diff --git a/test/runtests.jl b/test/runtests.jl
index b3cdc3a..35f65e4 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -247,4 +247,57 @@ end
 @enumx T=Typ FruitEmptyT
 @test instances(FruitEmptyT.Typ) == ()
 
+
+# Documented type (module) and instances
+begin
+    """
+    Documentation for FruitDoc
+    """
+    @enumx FruitDoc begin
+        "Apple documentation."
+        Apple
+        """
+        Banana documentation
+        on multiple lines.
+        """
+        Banana = 2
+        Orange = Apple
+    end
+    @eval const LINENUMBER = $(@__LINE__)
+    @eval const FILENAME = $(@__FILE__)
+    @eval const MODULE = $(@__MODULE__)
+end
+
+function get_doc_metadata(mod, s)
+    Base.Docs.meta(mod)[Base.Docs.Binding(mod, s)].docs[Union{}].data
+end
+
+@test FruitDoc.Apple === FruitDoc.T(0)
+@test FruitDoc.Banana === FruitDoc.T(2)
+@test FruitDoc.Orange === FruitDoc.T(0)
+
+mod_doc = @doc(FruitDoc)
+@test sprint(show, mod_doc) == "Documentation for FruitDoc\n"
+mod_doc_data = get_doc_metadata(FruitDoc, :FruitDoc)
+@test mod_doc_data[:linenumber] == LINENUMBER - 13
+@test mod_doc_data[:path] == FILENAME
+@test mod_doc_data[:module] == MODULE
+
+apple_doc = @doc(FruitDoc.Apple)
+@test sprint(show, apple_doc) == "Apple documentation.\n"
+apple_doc_data = get_doc_metadata(FruitDoc, :Apple)
+@test apple_doc_data[:linenumber] == LINENUMBER - 9
+@test apple_doc_data[:path] == FILENAME
+@test apple_doc_data[:module] == FruitDoc
+
+banana_doc = @doc(FruitDoc.Banana)
+@test sprint(show, banana_doc) == "Banana documentation on multiple lines.\n"
+banana_doc_data = get_doc_metadata(FruitDoc, :Banana)
+@test banana_doc_data[:linenumber] == LINENUMBER - 7
+@test banana_doc_data[:path] == FILENAME
+@test banana_doc_data[:module] == FruitDoc
+
+orange_doc = @doc(FruitDoc.Orange)
+@test startswith(sprint(show, orange_doc), "No documentation found")
+
 end # testset