From 7addf5fe01df85d91ee5fc7b4f7283bfa858e7fa Mon Sep 17 00:00:00 2001
From: jkrumbiegel <22495855+jkrumbiegel@users.noreply.github.com>
Date: Tue, 7 Dec 2021 10:47:39 +0100
Subject: [PATCH] fix docstring parsing issue (#41)

* fix docstring parsing issue

* bump patch version
---
 Project.toml     |  2 +-
 src/Chain.jl     | 44 +++++++++++++++++++++++++-------------------
 test/runtests.jl | 23 +++++++++++++++++++++++
 3 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/Project.toml b/Project.toml
index c3da82f..5cee706 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chain"
 uuid = "8be319e6-bccf-4806-a6f7-6fae938471bc"
 authors = ["Julius Krumbiegel"]
-version = "0.4.8"
+version = "0.4.9"
 
 [compat]
 julia = "1"
diff --git a/src/Chain.jl b/src/Chain.jl
index 19143ca..cc06f82 100644
--- a/src/Chain.jl
+++ b/src/Chain.jl
@@ -113,24 +113,8 @@ end
 rewrite(l::LineNumberNode, replacement) = (l, replacement)
 
 function rewrite_chain_block(firstpart, block)
-    block_expressions = block.args
-
-    # empty chain returns firstpart
-    if all(x -> x isa LineNumberNode, block_expressions)
-        return esc(firstpart)
-    end
-
-    rewritten_exprs = []
-    replacement = firstpart
-
-    for expr in block_expressions
-        rewritten, replacement = rewrite(expr, replacement)
-        push!(rewritten_exprs, rewritten)
-    end
-
-    result = Expr(:let, Expr(:block), Expr(:block, rewritten_exprs..., replacement))
-
-    :($(esc(result)))
+    pushfirst!(block.args, firstpart)
+    rewrite_chain_block(block)
 end
 
 """
@@ -191,6 +175,8 @@ function rewrite_chain_block(block)
     block_expressions = block.args
     isempty(block_expressions) && error("No expressions found in chain block.")
 
+    reconvert_docstrings!(block_expressions)
+
     # assign first line to first gensym variable
     firstvar = gensym()
     rewritten_exprs = []
@@ -198,7 +184,7 @@ function rewrite_chain_block(block)
 
     did_first = false
     for expr in block_expressions
-        # could be and expression first or a LineNumberNode, so a bit convoluted
+        # could be an expression first or a LineNumberNode, so a bit convoluted
         # we just do the firstvar transformation for the first non LineNumberNode
         # we encounter
         if !(did_first || expr isa LineNumberNode)
@@ -217,6 +203,26 @@ function rewrite_chain_block(block)
     :($(esc(result)))
 end
 
+# if a line in a chain is a string, it can be parsed as a docstring
+# for whatever is on the following line. because this is unexpected behavior
+# for most users, we convert all docstrings back to separate lines.
+function reconvert_docstrings!(args::Vector)
+    docstring_indices = findall(args) do arg
+        (arg isa Expr
+            && arg.head == :macrocall
+            && length(arg.args) == 4
+            && arg.args[1] == GlobalRef(Core, Symbol("@doc")))
+    end
+    # replace docstrings from back to front because this leaves the earlier indices intact
+    for i in reverse(docstring_indices)
+        e = args[i]
+        str = e.args[3]
+        nextline = e.args[4]
+        splice!(args, i:i, [str, nextline])
+    end
+    args
+end
+
 """
     @chain(block::Expr)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 86eb571..6bfdb52 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -412,4 +412,27 @@ end
         _ + 1
         @aside 1 + 2
     end
+end
+
+@testset "workaround for docstring parsing" begin
+    @test "hi" == @chain " hi " strip
+    @test "hi" == @chain " hi " begin
+        strip
+    end
+    @test "hi" == @chain begin
+        " hi "
+        strip
+    end
+    @test "hi" == @chain begin
+        "hi"
+        " $_ "
+        strip
+    end
+    @test "A" == @chain begin
+        'a'
+        " $_ "
+        strip
+        "$_"
+        uppercase
+    end
 end
\ No newline at end of file