Add arguments to @dropna, @replacena and @dissallowna

Author: davidanthoff

src/table_query_macros.jl

@@ -197,10 +197,29 @@ macro dissallowna()
     return :( Query.@map(map(our_get, _)) )
 end
 
+macro dissallowna(columns...)
+    return :( Query.@mutate( $( ( :( $(columns[i].value) = our_get(_.$(columns[i].value))  ) for i=1:length(columns) )... )   ) )
+end
+
 macro dropna()
     return :( i-> i |> Query.@filter(!any(isna, _)) |>  Query.@dissallowna() )
 end
 
-macro replacena(arg)
-    return :( Query.@map(map(i->our_get(i, $arg), _)) )
+macro dropna(columns...)
+    return :( i-> i |> Query.@filter(!any(($((:(isna(_.$(columns[i].value))) for i in 1:length(columns)  )...),))) |> Query.@dissallowna($(columns...)) )
+end
+
+macro replacena(arg, args...)
+    if length(args)==0 && !(arg isa Expr && arg.head==:call && length(arg.args)==3 && arg.args[1]==:(=>))
+        return :( Query.@map(map(i->our_get(i, $arg), _)) )
+    else
+        args = [arg; args...]
+
+        all(i isa Expr && i.head==:call && length(i.args)==3 && i.args[1]==:(=>)  for i in args) || error("Invalid syntax.")
+
+        columns = map(i->i.args[2].value, args)
+        replacement_values = map(i->i.args[3], args)
+
+        return :( Query.@mutate( $( ( :( $(columns[i]) = our_get(_.$(columns[i]), $(replacement_values[i]))  ) for i=1:length(columns) )... )   ) )
+    end
 end

test/test_macros.jl

@@ -57,6 +57,10 @@ end
     @test df |> @dropna() |> collect == [(a=1,b=1.), (a=3, b=3.)]
     @test df |> @filter(!any(isna, _)) |> @dropna() |> collect == [(a=1,b=1.), (a=3, b=3.)]
     @test df |> @select(:b) |> @dropna() |> collect == [(b=1.,),(b=2.,),(b=3.,)]
+
+    @test df |> @dropna(:a) |> collect == [(a=1,b=1.), (a=3, b=3.)]
+    @test df |> @dropna(:b) |> collect == [(a=DataValue(1),b=1.), (a=DataValue{Int}(),b=2.),(a=DataValue(3), b=3.)]
+    @test df |> @dropna(:a, :b) |> collect == [(a=1,b=1.), (a=3, b=3.)]
 end
 
 @testset "@replacena" begin
@@ -66,12 +70,18 @@ end
     @test df |> @replacena(2) |> collect == [(a=1,b=1.), (a=2, b=2.), (a=3, b=3.)]
     @test df |> @dropna() |> @replacena(2) |> collect == [(a=1,b=1.), (a=3, b=3.)]
     @test df |> @select(:b) |> @replacena(2) |> collect == [(b=1.,),(b=2.,),(b=3.,)]
+
+    @test df |> @replacena(:a=>2) |> collect == [(a=1,b=1.), (a=2, b=2.), (a=3, b=3.)]
+    @test df |> @replacena(:b=>2) |> collect == [(a=DataValue(1),b=1.), (a=DataValue{Int}(),b=2.),(a=DataValue(3), b=3.)]
+    @test df |> @replacena(:a=>2, :b=>8) |> collect == [(a=1,b=1.), (a=2, b=2.), (a=3, b=3.)]
 end
 
 @testset "@dissallowna" begin
 
     df = DataFrame(a=[1,missing,3], b=[1.,2.,3.])
 
     @test_throws DataValueException df |> @dissallowna() |> collect
-    df |> @filter(!any(isna, _)) |> @dissallowna() |> collect == [(a=1,b=1.), (a=3, b=3.)]
+    @test df |> @filter(!any(isna, _)) |> @dissallowna() |> collect == [(a=1,b=1.), (a=3, b=3.)]
+    @test_throws DataValueException df |> @dissallowna(:a) |> collect
+    @test df |> @dissallowna(:b) |> collect == [(a=DataValue(1),b=1.), (a=DataValue{Int}(),b=2.),(a=DataValue(3), b=3.)]
 end