fix nested single line chain (#43)

* fix nested single line chain * bump patch version
jkrumbiegel · Dec 8, 2021 · 5bfc333 · 5bfc333 · jkrumbiegel · Dec 8, 2021
1 parent 7addf5f
commit 5bfc333
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chain"
 uuid = "8be319e6-bccf-4806-a6f7-6fae938471bc"
 authors = ["Julius Krumbiegel"]
-version = "0.4.9"
+version = "0.4.10"
 
 [compat]
 julia = "1"

diff --git a/src/Chain.jl b/src/Chain.jl
@@ -257,15 +257,15 @@ function replace_underscores(expr::Expr, replacement)
     # if a @chain macrocall is found, only its first arg can be replaced if it's an
     # underscore, otherwise the macro insides are left untouched
     if expr.head == :macrocall && expr.args[1] == Symbol("@chain")
-        length(expr.args) != 4 && error("Malformed nested @chain macro")
+        length(expr.args) < 3 && error("Malformed nested @chain macro")
         expr.args[2] isa LineNumberNode || error("Malformed nested @chain macro")
         arg3 = if expr.args[3] == Symbol("_")
             found_underscore = true
             replacement
         else
             expr.args[3]
         end
-        newexpr = Expr(:macrocall, Symbol("@chain"), expr.args[2], arg3, expr.args[4])
+        newexpr = Expr(:macrocall, Symbol("@chain"), expr.args[2], arg3, expr.args[4:end]...)
     # for all other expressions, their arguments are checked for underscores recursively
     # and replaced if any are found
     else

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -435,4 +435,10 @@ end
         "$_"
         uppercase
     end
+end
+
+@testset "nested single line chain" begin
+    @test 36 == @chain 1:3 begin
+        @chain _ sum _ ^ 2
+    end
 end