add Not to delete!

JuliaData · Apr 27, 2020 · 3c5013a · 3c5013a
1 parent b1f675d
commit 3c5013a
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 57 deletions.
diff --git a/src/dataframe/dataframe.jl b/src/dataframe/dataframe.jl
@@ -836,8 +836,8 @@ end
 
 Delete rows specified by `inds` from a `DataFrame` `df` in place and return it.
 
-Internally `deleteat!` is called for all columns so `inds` must
-be: a vector of sorted and unique integers, a boolean vector or an integer.
+Internally `deleteat!` is called for all columns so `inds` must be:
+a vector of sorted and unique integers, a boolean vector, an integer, or `Not`.
 
 # Examples
 ```jldoctest
@@ -865,8 +865,9 @@ function Base.delete!(df::DataFrame, inds)
         throw(BoundsError(df, (inds, :)))
     end
     # we require ind to be stored and unique like in Base
+    # otherwise an error will be thrown and the data frame will get corrupted
     foreach(col -> deleteat!(col, inds), _columns(df))
-    df
+    return df
 end
 
 function Base.delete!(df::DataFrame, inds::AbstractVector{Bool})
@@ -875,9 +876,11 @@ function Base.delete!(df::DataFrame, inds::AbstractVector{Bool})
     end
     drop = findall(inds)
     foreach(col -> deleteat!(col, drop), _columns(df))
-    df
+    return df
 end
 
+Base.delete!(df::DataFrame, inds::Not) = delete!(df, axes(df, 1)[inds])
+
 ##############################################################################
 ##
 ## Hcat specialization

diff --git a/test/dataframe.jl b/test/dataframe.jl
@@ -549,14 +549,6 @@ end
     @test delete!(df, 2) === df
     @test df == DataFrame(a=[1], b=[3.0])
 
-    df = DataFrame(a=[1, 2, 3], b=[3.0, 4.0, 5.0])
-    @test delete!(df, 2:3) === df
-    @test df == DataFrame(a=[1], b=[3.0])
-
-    df = DataFrame(a=[1, 2, 3], b=[3.0, 4.0, 5.0])
-    @test delete!(df, [2, 3]) === df
-    @test df == DataFrame(a=[1], b=[3.0])
-
     df = DataFrame(a=Union{Int, Missing}[1, 2], b=Union{Float64, Missing}[3.0, 4.0])
     @test delete!(df, 1) === df
     @test df == DataFrame(a=[2], b=[4.0])
@@ -565,13 +557,15 @@ end
     @test delete!(df, 2) === df
     @test df == DataFrame(a=[1], b=[3.0])
 
-    df = DataFrame(a=Union{Int, Missing}[1, 2, 3], b=Union{Float64, Missing}[3.0, 4.0, 5.0])
-    @test delete!(df, 2:3) === df
-    @test df == DataFrame(a=[1], b=[3.0])
+    for v in (2:3, [2, 3])
+        df = DataFrame(a=Union{Int, Missing}[1, 2, 3], b=Union{Float64, Missing}[3.0, 4.0, 5.0])
+        @test delete!(df, v) === df
+        @test df == DataFrame(a=[1], b=[3.0])
 
-    df = DataFrame(a=Union{Int, Missing}[1, 2, 3], b=Union{Float64, Missing}[3.0, 4.0, 5.0])
-    @test delete!(df, [2, 3]) === df
-    @test df == DataFrame(a=[1], b=[3.0])
+        df = DataFrame(a=[1, 2, 3], b=[3.0, 4.0, 5.0])
+        @test delete!(df, v) === df
+        @test df == DataFrame(a=[1], b=[3.0])
+    end
 
     df = DataFrame()
     @test_throws BoundsError delete!(df, 10)
@@ -589,45 +583,19 @@ end
     @test delete!(df, [false, true, false]) === df
     @test df == DataFrame(a=[1, 3], b=[3, 1])
 
-    x = [1, 2, 3]
-    df = DataFrame(x=x)
-    @test delete!(df, 1) == DataFrame(x=[2, 3])
-    @test x == [1, 2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x)
-    @test delete!(df, [1]) == DataFrame(x=[2, 3])
-    @test x == [1, 2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x)
-    @test delete!(df, 1:1) == DataFrame(x=[2, 3])
-    @test x == [1, 2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x)
-    @test delete!(df, [true, false, false]) == DataFrame(x=[2, 3])
-    @test x == [1, 2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x, copycols=false)
-    @test delete!(df, 1) == DataFrame(x=[2, 3])
-    @test x == [2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x, copycols=false)
-    @test delete!(df, [1]) == DataFrame(x=[2, 3])
-    @test x == [2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x, copycols=false)
-    @test delete!(df, 1:1) == DataFrame(x=[2, 3])
-    @test x == [2, 3]
-
-    x = [1, 2, 3]
-    df = DataFrame(x=x, copycols=false)
-    @test delete!(df, [true, false, false]) == DataFrame(x=[2, 3])
-    @test x == [2, 3]
+    for v in (1, [1], 1:1, [true, false, false])
+        x = [1, 2, 3]
+        df = DataFrame(x=x)
+        @test delete!(df, v) == DataFrame(x=[2, 3])
+        @test x == [1, 2, 3]
+    end
+
+    for v in (1, [1], 1:1, [true, false, false], Not(2,3), Not([false, true, true]))
+        x = [1, 2, 3]
+        df = DataFrame(x=x, copycols=false)
+        @test delete!(df, v) == DataFrame(x=[2, 3])
+        @test x == [2, 3]
+    end
 end
 
 @testset "describe" begin