Merge pull request #258 from pdeffebach/linenums_fix

Add dependency on MacroTools to avoid `Base.remove_linenums!`

Project.toml

@@ -4,12 +4,14 @@ version = "0.7.0"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [compat]
 DataFrames = "0.21, 0.22, 1"
 Reexport = "0.2, 1"
 julia = "1"
+MacroTools = "0.5"
 
 [extras]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

src/DataFramesMeta.jl

@@ -2,6 +2,8 @@ module DataFramesMeta
 
 using Reexport
 
+using MacroTools
+
 @reexport using DataFrames
 
 # Basics:

src/parsing.jl

@@ -118,17 +118,14 @@ julia> MacroTools.prettify(fun)
 
 """
 function get_source_fun(function_expr; wrap_byrow::Bool=false)
+    function_expr = MacroTools.unblock(function_expr)
+
     # recursive step for begin :a + :b end
     if is_macro_head(function_expr, "@byrow")
         if wrap_byrow
             throw(ArgumentError("Redundant `@byrow` calls."))
         end
         return get_source_fun(function_expr.args[3], wrap_byrow=true)
-    elseif function_expr isa Expr &&
-        function_expr.head == :block &&
-        length(function_expr.args) == 1
-
-        return get_source_fun(function_expr.args[1])
     elseif is_simple_non_broadcast_call(function_expr)
         source = args_to_selectors(function_expr.args[2:end])
         fun_t = function_expr.args[1]
@@ -230,13 +227,7 @@ function fun_to_vec(ex::Expr; nolhs::Bool = false, gensym_names::Bool = false, w
 
     if !nokw
         lhs = ex.args[1]
-        rhs_t = ex.args[2]
-        # if lhs is a cols(y) then the rhs gets parsed as a block
-        if onearg(lhs, :cols) && rhs_t.head === :block && length(rhs_t.args) == 1
-            rhs = rhs_t.args[1]
-        else
-            rhs = rhs_t
-        end
+        rhs = MacroTools.unblock(ex.args[2])
     else
         throw(ArgumentError("This path should not be reached"))
     end
@@ -281,7 +272,6 @@ function fun_to_vec(ex::Expr; nolhs::Bool = false, gensym_names::Bool = false, w
     if onearg(lhs, :cols) && onearg(rhs, :cols)
         source = rhs.args[2]
         dest = lhs.args[2]
-
         return quote
             $source => $dest
         end
@@ -347,7 +337,7 @@ of expression-like object (`Expr`, `QuoteNode`, etc.),
 puts them into a single array, removing line numbers.
 """
 function create_args_vector(args...)
-    Any[Base.remove_linenums!(arg) for arg in args], false
+    create_args_vector(Expr(:block, args...))
 end
 
 """
@@ -380,9 +370,9 @@ function create_args_vector(arg)
     end
 
     if arg isa Expr && arg.head == :block
-        x = Base.remove_linenums!(arg).args
+        x = MacroTools.rmlines(arg).args
     else
-        x = Any[Base.remove_linenums!(arg)]
+        x = Any[arg]
     end
 
     if wrap_byrow && any(t -> is_macro_head(t, "@byrow"), x)

test/dataframes.jl

@@ -728,4 +728,47 @@ end
     @test @transform(df, cols("X Y Z")) == df
 end
 
+macro linenums_macro(arg)
+    if arg isa Expr && arg.head == :block && length(arg.args) == 1 && arg.args[1] isa LineNumberNode
+        esc(:(true))
+    else
+        esc(:(false))
+    end
+end
+
+@testset "removing lines" begin
+    df = DataFrame(a = [1], b = [2])
+    # Can't use @test because @test remove line numbers
+    d = @transform(df, y = @linenums_macro begin end)
+    @test d.y == [true]
+
+    d = @transform df begin
+        y = @linenums_macro begin end
+    end
+
+    @test d.y == [true]
+
+    d = @transform df @byrow begin
+        y = @linenums_macro begin end
+    end
+
+    @test d.y == [true]
+
+    d = @where(df, @linenums_macro begin end)
+
+    @test nrow(d) == 1
+
+    d = @where df begin
+        @byrow @linenums_macro begin end
+    end
+
+    @test nrow(d) == 1
+
+    d = @where df @byrow begin
+        @linenums_macro begin end
+    end
+
+    @test nrow(d) == 1
+end
+
 end # module