From d93dcbd48fa988e43176fd41bffe823d26c379dd Mon Sep 17 00:00:00 2001
From: pdeffebach <23196228+pdeffebach@users.noreply.github.com>
Date: Thu, 14 Nov 2024 17:04:02 -0500
Subject: [PATCH] Fix multi-column select! #404 (#405)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* add allow multicol

* add tests

* add fix

* Update src/macros.jl

Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl>

---------

Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl>
---
 src/macros.jl    |  6 ++---
 test/multicol.jl | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/macros.jl b/src/macros.jl
index fd653549..2fb2dbd1 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -1930,7 +1930,7 @@ end
 function rselect_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = true)
 
-    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select($x, $(t...); $(kw...))
     end
@@ -1983,7 +1983,7 @@ end
 function select!_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = false)
 
-    t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select!($x, $(t...); $(kw...))
     end
@@ -2119,7 +2119,7 @@ end
 function rselect!_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = true)
 
-    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select!($x, $(t...); $(kw...))
     end
diff --git a/test/multicol.jl b/test/multicol.jl
index 2e44607c..0caf6444 100644
--- a/test/multicol.jl
+++ b/test/multicol.jl
@@ -22,6 +22,67 @@ df = DataFrame(A = 1, AA = 2, B = 3)
     @test t == DataFrame(AA = 2, B = 3)
 end
 
+@testset "rselect_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    t = @rselect df Not(:A)
+    @test t == DataFrame(AA = 2, B = 3)
+
+    t = @rselect df All()
+    @test t == DataFrame(A = 1, AA = 2, B = 3)
+
+    t = @rselect df Cols(r"A")
+    @test t == DataFrame(A = 1, AA = 2)
+
+    t = @rselect df Between(:AA, :B)
+    @test t == DataFrame(AA = 2, B = 3)
+end
+
+@testset "select!_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Not(:A)
+    @test df == DataFrame(AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df All()
+    @test df == DataFrame(A = 1, AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Cols(r"A")
+    @test df == DataFrame(A = 1, AA = 2)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Between(:AA, :B)
+    @test df == DataFrame(AA = 2, B = 3)
+end
+
+@testset "rselect!_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Not(:A)
+    @test df == DataFrame(AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df All()
+    @test df == DataFrame(A = 1, AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Cols(r"A")
+    @test df == DataFrame(A = 1, AA = 2)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Between(:AA, :B)
+    @test df == DataFrame(AA = 2, B = 3)
+end
+
+
 @testset "othermacros_multi" begin
     df = DataFrame(A = 1, AA = 2, B = 3)