switch to MacroTools

Author: Nathan Smith

README.md

@@ -5,15 +5,13 @@ Python generators are great, Julia is great. Why compromise?
 ```julia
 julia> using PyGen
 
-julia> @pygen """
-       function fibonacci()
+julia> @pygen function fibonacci()
            n, m = 0, 1
            while true
-               yield m
+               yield(m)
                n, m = m, n + m
            end 
        end
-       """
 fibonacci (generic function with 1 method)
 
 julia> for i in fibonacci()
@@ -32,4 +30,4 @@ julia> for i in fibonacci()
 .
 ```
 
-Voila! Generators without all the ceremony!
+Voila! Generators without all the ceremony!

REQUIRE

@@ -1 +1,2 @@
 julia 0.5
+MacroTools 0.3.0

src/PyGen.jl

@@ -1,57 +1,71 @@
 module PyGen
 
-export @pygen
-#=
-    @pygen is a macro that allows for python like generator functions.
-The idea is that a function like this:
-
-function f(args)
-    ...
-    yield x
-    ...
-    yield y
-    ...
-end
-
-should become....
-
-function f(args)
-    function γ(c::Channel)
-        ...
-        put!(c, x)
-        ...
-        put!(c, y)
-        ...
-    end
-    return Channel(γ)
-end
-=#
-
-const yregex = r"(yield .*)"
-
-
-function asput(str)
-    # Convert "yield λ" to "put!(c, λ)
-    λ = replace(str, "yield ", "")
-    return "put!(c, $λ)"
-end
-
+using MacroTools
 
+export @pygen
 """
     @pygen
 
 Making julie walking a little *more* like python. Stick 
 the @pygen macro in front of a function declaration and
 you can use `yield` statements to make a python style
-generator! 
+generator!
+
+## Examples
+
+```julia
+julia> @pygen function fibonacci()
+           n, m = 0, 1
+           while true
+               yield(m)
+               n, m = m, n + m
+           end
+       end
+fibonacci (generic function with 1 method)
+
+julia> for n in fibonacci()
+           println(n)
+           sleep(1)
+       end
+1
+1
+2
+3
+5
+8
+13
+.
+.
+.
+``` 
 """
 macro pygen(f)
-    f′ = replace(f, yregex, asput)
-    lines = f′ |> IOBuffer |> readlines
-    ind = findfirst(x -> contains(x, "function"), lines)    
-    insert!(lines, ind + 1, "function γ(c::Channel)")
-    push!(lines, "return Channel(γ)\nend")
-    return join(lines, "\n") |> parse |> esc 
+
+    # generate a new symbol for the channel name and 
+    # wrapper function
+    c = gensym()
+    η = gensym()
+
+    # yield(λ) → put!(c, λ)
+    f′ = MacroTools.postwalk(f) do x
+        @capture(x, yield(λ_)) || return x
+        return :(put!($c, $λ))
+    end 
+    
+    # Fetch the function name and args
+    @capture(f′, function func_(args__) body_ end)
+
+    # wrap up the η function
+    final = quote 
+        function $func($(args...))
+            function $η($c)
+                $body
+            end
+            return Channel(c -> $η(c))
+        end
+    end
+
+    return esc(final)
 end
 
 end 

test/runtests.jl

@@ -2,36 +2,27 @@ using PyGen
 using Base.Test
 
 # write your own tests here
-@test PyGen.asput("yield pi") == "put!(c, pi)"
-
-@pygen "
-
-function λ()
-    yield 10
+@pygen function λ()
+    yield(10)
 end
-"
 
 # Test for \n in first line issue
 @test isdefined(:λ) == true
 
 # Test that the poodle function doesn't error out
-@pygen "
-function 🐩()
-    yield \"🐩\"
+@pygen function 🐩()
+    yield("🐩")
 end
-"
 
 @test isdefined(:🐩) == true
 
 # Test that one of these generators actually works
-@pygen """
-function squares(n)
+@pygen function squares(n)
     i = 0
     while i <= n
-        yield i^2
+        yield(i^2)
         i += 1
     end
 end
-"""
 
 @test sum(squares(10)) == sum(i^2 for i in 1:10)