From c59cbb2ccc0d61f71f815e3cbdf6a5c5e3ac603a Mon Sep 17 00:00:00 2001
From: recloser <44084068+recloser@users.noreply.github.com>
Date: Wed, 15 Jan 2020 11:48:30 +0100
Subject: [PATCH] add support for fields defined in cases

---
 tests/tests.nim | 27 ++++++++++++++++++++++++++-
 with.nim        |  9 ++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/tests/tests.nim b/tests/tests.nim
index efb932d..aae4063 100644
--- a/tests/tests.nim
+++ b/tests/tests.nim
@@ -17,6 +17,22 @@ type
   Sub = object
     foo: Foo
 
+  Cases = object
+    x: int
+    case cond: int8
+    of 1:
+      case secondCond: bool
+      of true: 
+        y: int
+      of false: 
+        z: string
+        w: int
+    of 2:
+      a: int
+    else:
+      b: float
+
+
 suite "with":
 
   test "basic object":
@@ -161,4 +177,13 @@ suite "with":
       check field1 == 100
       check field2 == "something"
 
-
+  test "fields in case":
+    var cases = Cases(cond: 1, secondCond: false, x: 123, z: "hello", w: 321)
+    with cases:
+      check secondCond == false
+      check x == 123
+      check z == "hello"
+      check w == 321
+      check compiles(y == 0)
+      check compiles(a == 0)
+      check compiles(b == 0.0)
diff --git a/with.nim b/with.nim
index 931284b..1b408fe 100755
--- a/with.nim
+++ b/with.nim
@@ -7,6 +7,13 @@
 
 import macros
 
+proc getAllIdentDefs(x: NimNode): seq[NimNode] =
+  for n in x:
+    if n.kind == nnkIdentDefs:
+      result &= n
+    elif n.kind in {nnkRecCase, nnkRecList, nnkOfBranch, nnkElse}:
+      result &= getAllIdentDefs(n)
+
 proc createInner(x: NimNode): NimNode =
   result = newStmtList()
   var t = getTypeImpl(x)
@@ -14,7 +21,7 @@ proc createInner(x: NimNode): NimNode =
     t = getTypeImpl(t[0])
   assert(t.kind in {nnkObjectTy, nnkTupleTy}, "`with` must be called with " &
     "either objects or tuples.")
-  for field in (if t.kind == nnkObjectTy: t[2] else: t):
+  for field in getAllIdentDefs(if t.kind == nnkObjectTy: t[2] else: t):
     let name = newIdentNode($field[0])
     if field[1].kind == nnkProcTy:
       result.add quote do: