From eb07a5a75b63110642c5ce6f9126c9c8af231a64 Mon Sep 17 00:00:00 2001
From: Timothee Cour <timothee.cour2@gmail.com>
Date: Wed, 10 Mar 2021 08:08:24 -0800
Subject: [PATCH] add typetraits.OrdinalEnum, enumutils.symbolName (#17281)

---
 changelog.md                |  3 ++-
 lib/pure/typetraits.nim     | 14 +++++++++-----
 lib/std/enumutils.nim       | 28 +++++++++++++++++++++++++---
 tests/stdlib/tenumutils.nim | 21 +++++++++++++++++++++
 4 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/changelog.md b/changelog.md
index 67c37589f0ae..05211b11663e 100644
--- a/changelog.md
+++ b/changelog.md
@@ -47,8 +47,9 @@
 
 - Added `std/enumutils` module. Added `genEnumCaseStmt` macro that generates case statement to parse string to enum.
   Added `items` for enums with holes.
+  Added `symbolName` to return the enum symbol name ignoring the human readable name.
 
-- Added `typetraits.SomeEnumWithHoles` for enums with holes.
+- Added `typetraits.HoleyEnum` for enums with holes, `OrdinalEnum` for enums without holes.
 
 - Removed deprecated `iup` module from stdlib, it has already moved to
   [nimble](https://github.com/nim-lang/iup).
diff --git a/lib/pure/typetraits.nim b/lib/pure/typetraits.nim
index 69c0329ef1a5..6827e23b1a2e 100644
--- a/lib/pure/typetraits.nim
+++ b/lib/pure/typetraits.nim
@@ -15,16 +15,20 @@
 import std/private/since
 export system.`$` # for backward compatibility
 
-type SomeEnumWithHoles* = (not Ordinal) and enum ## Enum with holes.
+type HoleyEnum* = (not Ordinal) and enum ## Enum with holes.
+type OrdinalEnum* = Ordinal and enum ## Enum without holes.
 
 runnableExamples:
   type A = enum a0 = 2, a1 = 4, a2
   type B = enum b0 = 2, b1, b2
-  assert A is SomeEnumWithHoles
-  assert B isnot SomeEnumWithHoles
-  assert int isnot SomeEnumWithHoles
+  assert A is enum
+  assert A is HoleyEnum
+  assert A isnot OrdinalEnum
+  assert B isnot HoleyEnum
+  assert B is OrdinalEnum
+  assert int isnot HoleyEnum
   type C[T] = enum h0 = 2, h1 = 4
-  assert C[float] is SomeEnumWithHoles
+  assert C[float] is HoleyEnum
 
 proc name*(t: typedesc): string {.magic: "TypeTrait".} =
   ## Returns the name of the given type.
diff --git a/lib/std/enumutils.nim b/lib/std/enumutils.nim
index 56a6d82a7e7d..16dab9d1aabb 100644
--- a/lib/std/enumutils.nim
+++ b/lib/std/enumutils.nim
@@ -8,6 +8,7 @@
 #
 
 import std/macros
+from std/typetraits import OrdinalEnum, HoleyEnum
 
 # xxx `genEnumCaseStmt` needs tests and runnableExamples
 
@@ -65,10 +66,17 @@ macro genEnumCaseStmt*(typ: typedesc, argSym: typed, default: typed,
     expectKind(default, nnkSym)
     result.add nnkElse.newTree(default)
 
-macro enumWithHolesFullRange(a: typed): untyped = 
+macro enumFullRange(a: typed): untyped =
   newNimNode(nnkCurly).add(a.getType[1][1..^1])
 
-iterator items*[T: enum and not Ordinal](E: typedesc[T]): T =
+macro enumNames(a: typed): untyped =
+  # this could be exported too; in particular this could be useful for enum with holes.
+  result = newNimNode(nnkBracket)
+  for ai in a.getType[1][1..^1]:
+    assert ai.kind == nnkSym
+    result.add newLit ai.strVal
+
+iterator items*[T: HoleyEnum](E: typedesc[T]): T =
   ## Iterates over an enum with holes.
   runnableExamples:
     type A = enum a0 = 2, a1 = 4, a2
@@ -76,4 +84,18 @@ iterator items*[T: enum and not Ordinal](E: typedesc[T]): T =
     from std/sequtils import toSeq
     assert A.toSeq == [a0, a1, a2]
     assert B[float].toSeq == [B[float].b0, B[float].b1]
-  for a in enumWithHolesFullRange(E): yield a
+  for a in enumFullRange(E): yield a
+
+func symbolName*[T: OrdinalEnum](a: T): string =
+  ## Returns the symbol name of an enum.
+  runnableExamples:
+    type B = enum
+      b0 = (10, "kb0")
+      b1 = "kb1"
+      b2
+    let b = B.low
+    assert b.symbolName == "b0"
+    assert $b == "kb0"
+    static: assert B.high.symbolName == "b2"
+  const names = enumNames(T)
+  names[a.ord - T.low.ord]
diff --git a/tests/stdlib/tenumutils.nim b/tests/stdlib/tenumutils.nim
index dd5da1974856..11142216c422 100644
--- a/tests/stdlib/tenumutils.nim
+++ b/tests/stdlib/tenumutils.nim
@@ -12,5 +12,26 @@ template main =
     doAssert A.toSeq == [a0, a1, a2]
     doAssert B[float].toSeq == [B[float].b0, B[float].b1]
 
+  block: # symbolName
+    block:
+      type A2 = enum a20, a21, a22
+      doAssert $a21 == "a21"
+      doAssert a21.symbolName == "a21"
+      proc `$`(a: A2): string = "foo"
+      doAssert $a21 == "foo"
+      doAssert a21.symbolName == "a21"
+      var a = a22
+      doAssert $a == "foo"
+      doAssert a.symbolName == "a22"
+
+    type B = enum
+      b0 = (10, "kb0")
+      b1 = "kb1"
+      b2
+    let b = B.low
+    doAssert b.symbolName == "b0"
+    doAssert $b == "kb0"
+    static: doAssert B.high.symbolName == "b2"
+
 static: main()
 main()