Add enumerate extension

README.md

@@ -9,7 +9,7 @@ Extensions so far include:
 
   This module was modified from the OOP section on the excellent website http://nim-by-example.github.io/
 
-  usage: 
+  Usage: 
   
   ```nim
   import extensions/oop
@@ -35,3 +35,21 @@ Extensions so far include:
   ```bash
   nim c -r tests/test1.nim
   ```
+
+## Enumerate
+
+    This macro is useful for enumerating on any iterator. It will autodetect which iterator you want to use and give you a useful enumeration.
+
+  Usage:
+
+  ```nim
+  import os
+  import extensions/enumerate
+
+  for i, line in enumerate(lines, "file.txt"):
+      echo i, " ", line
+
+  for i, tup in enumerate(walkDir, "."):
+      echo i, " ", tup
+  ```
+

extensions.nimble

@@ -1,6 +1,6 @@
 # Package
 
-version       = "0.2.0"
+version       = "0.2.1"
 author        = "Joey Yakimowich-Payne"
 description   = "Nim extensions"
 license       = "MIT"
@@ -8,4 +8,4 @@ srcDir        = "src"
 
 # Dependencies
 
-requires "nim >= 0.18.0"
+requires "nim >= 0.18.1"

src/extensions/enumerate.nim

@@ -0,0 +1,75 @@
+import macros
+import extensions/nimnode
+
+proc argsMatch(impl, args: NimNode): bool =
+  var implArgs = impl["FormalParams"]
+  result = true
+  for i, arg in args.pairs:
+    let implIndex = i + 1 # skip the return type
+    let symType = arg.getType()
+    let implType = implArgs[implIndex][1]
+    result = result and ($symType == $implType)
+
+macro enumerate*(iter: typed, args: varargs[typed]): untyped =
+  ## Takes in an i
+  var iterDef: NimNode
+
+  if len(iter) > 1:
+    # If there are multiple syms to choose from,
+    # choose the one with matching args
+    for c in iter:
+      let impl = c.getImpl().copy()
+      if argsMatch(impl, args):
+        iterDef = impl
+  else:
+    iterDef = iter.getImpl().copy()
+
+  # set the symbol to empty because
+  # we don't need it
+  iterDef["Sym"] = newEmptyNode()
+
+  # remove all stmts, we only want the def
+  # for now
+  iterDef["StmtList"] = newEmptyNode()
+
+  # remove pragma
+  iterDef["Pragma"] = newEmptyNode()
+
+  var
+    returnType = iterDef["FormalParams > 0"]
+    iterCall = newNimNode(nnkCall).add(iter.copy())
+
+  for child in iterDef[3].children:
+    if child.kind == nnkIdentDefs:
+      iterCall.add(child[0].copy())
+
+  # mod the return type to return both int and original type
+  iterDef["FormalParams > 0"] = newNimNode(nnkPar).add(ident"int", returnType)
+  #iterDef[3][0] = newNimNode(nnkPar).add(ident"int", returnType)
+
+  template iterBody(iterCall) =
+    var i = 0
+    for tup in iterCall:
+      # Since Nim's checker won't allow a yield statement here,
+      # we need to get creative
+      replace(i, tup)
+      inc i
+
+  var stmtList = getAst(iterBody(iterCall))
+
+  var replaceCall = stmtList["ForStmt > StmtList > Call"]
+
+  var yieldStmt = nnkYieldStmt.newTree(
+    nnkPar.newTree(
+      replaceCall[1],
+      replaceCall[2]
+    )
+  )
+
+  stmtList["ForStmt > StmtList > Call"] = yieldStmt
+  iterDef[^1] = stmtList
+
+  # Call the iterator surrounded by parens
+  result = nnkCall.newTree(nnkPar.newTree(iterDef))
+  for arg in args:
+    result.add(arg)

src/extensions/nim.cfg

@@ -0,0 +1 @@
+--path="../"

src/extensions/nimnode.nim

@@ -0,0 +1,85 @@
+import macros, strutils
+
+proc `[]=`*(node: NimNode, nodePath: string, value: NimNode) =
+  ## Set a NimNode based on a string search path like "StmtList > Command > Sym"
+  let nodeNames = nodePath.replace(" ", "").split(">")
+
+  let targetIndex = len(nodeNames) - 1
+
+  var
+    index = -1
+    targetParent: NimNode
+    stack: seq[seq[NimNode]] = @[@[node]]
+    indexList: seq[int] = newSeq[int](len(nodeNames))
+
+
+  while stack.len() > 0 and index < targetIndex:
+    let parentList = stack.pop()
+    var newParents: seq[NimNode] = @[]
+
+    block WhileBlock:
+      for parent in parentList:
+        inc index
+        for chIndex, child in parent.pairs:
+          let kindStr = ($child.kind)[3..^1]
+
+          var tryParse = -1
+          try:
+            tryParse = nodeNames[index].parseInt
+          except ValueError:
+            discard
+
+          if kindStr == nodeNames[index] or tryParse == chIndex:
+            targetParent = parent
+            newParents.add(child)
+            indexList[index] = chIndex
+
+            if index == targetIndex:
+              break WhileBlock
+
+    if newParents.len() > 0:
+      stack.add(newParents)
+
+  if index == targetIndex:
+    let modIndex = indexList.pop()
+    var curNode = node
+    for i in indexList:
+      curNode = curNode[i]
+
+    curNode[modIndex] = value
+
+
+proc `[]`*(node: NimNode, nodePath: string): NimNode =
+  ## Get a NimNode based on a string search path like "StmtList > Command > Sym"
+  let nodeNames = nodePath.replace(" ", "").split(">")
+
+  let targetIndex = len(nodeNames) - 1
+
+  var
+    index = -1
+    stack: seq[seq[NimNode]] = @[@[node]]
+
+  while stack.len() > 0 and index < targetIndex:
+    let parentList = stack.pop()
+    var newParents: seq[NimNode] = @[]
+
+    block WhileBlock:
+      for parent in parentList:
+        inc index
+        for i, child in parent.pairs:
+          let kindStr = ($child.kind)[3..^1]
+
+          var tryParse = -1
+          try:
+            tryParse = nodeNames[index].parseInt
+          except ValueError:
+            discard
+
+          if kindStr == nodeNames[index] or tryParse == i:
+            result = child
+            newParents.add(child)
+            if index == targetIndex:
+              break WhileBlock
+
+    if newParents.len() > 0:
+      stack.add(newParents)