add haskey and get support for DataFrameColumns (#3282)

Author: bkamins

NEWS.md

@@ -23,6 +23,9 @@
 * Add `keep` keyword argument to `nonunique`, `unique`, and `unique!`
   allowing to specify which duplicate rows should be kept
   ([#3260](https://github.com/JuliaData/DataFrames.jl/pull/3260))
+* Add `haskey` and `get` methods to `DataFrameColumns`
+  to make it support dictionary interface more completely
+  ([#3282](https://github.com/JuliaData/DataFrames.jl/pull/3282))
 
 ## Bug fixes
 

src/abstractdataframe/iteration.jl

@@ -240,6 +240,14 @@ with the corresponding column vector, i.e. `name => col`
 where `name` is the column name of the column `col`.
 """
 Base.pairs(itr::DataFrameColumns) = Base.Iterators.Pairs(itr, keys(itr))
+
+Base.haskey(itr::DataFrameColumns, col::Union{AbstractString, Symbol}) =
+    columnindex(parent(itr), col) > 0
+Base.haskey(itr::DataFrameColumns, col::Union{Signed, Unsigned}) =
+    0 < col <= ncol(parent(itr))
+Base.get(itr::DataFrameColumns, col::ColumnIndex, default) =
+    haskey(itr, col) ? itr[col] : default
+
 Base.findnext(f::Function, itr::DataFrameColumns, i::Integer) =
     findnext(f, values(itr), i)
 Base.findnext(f::Function, itr::DataFrameColumns, i::Union{Symbol, AbstractString}) =

test/iteration.jl

@@ -247,6 +247,39 @@ end
         @test parent(er3) isa SubDataFrame
         @test parentindices(parent(er3))[1] == 1:2
     end
+
+end
+
+@testset "haskey and get for DataFrameColumns" begin
+    df_ref = DataFrame(a=1:3, b=2:4, c=3:5)
+    for df in (df_ref, @view df_ref[1:3, 1:2])
+        dfc = eachcol(df)
+        @test !haskey(dfc, 0)
+        @test haskey(dfc, 1)
+        @test haskey(dfc, 2)
+        @test !haskey(dfc, 4)
+        @test !haskey(dfc, 0x0)
+        @test haskey(dfc, 0x1)
+        @test !haskey(dfc, :x)
+        @test !haskey(dfc, "x")
+        @test !haskey(dfc, Test.GenericString("x"))
+        @test haskey(dfc, :a)
+        @test haskey(dfc, "a")
+        @test haskey(dfc, Test.GenericString("a"))
+
+        @test get(dfc, 0, "error") == "error"
+        @test get(dfc, 1, "error") == 1:3
+        @test get(dfc, 4, "error") == "error"
+        @test get(dfc, 0x0, "error") == "error"
+        @test get(dfc, 0x1, "error") == 1:3
+        @test get(dfc, :x, "error") == "error"
+        @test get(dfc, "x", "error") == "error"
+        @test get(dfc, Test.GenericString("x"), "error") == "error"
+        @test get(dfc, :a, "error") == 1:3
+        @test get(dfc, "a", "error") == 1:3
+        @test get(dfc, Test.GenericString("a"), "error") == 1:3
+        @test_throws MethodError get(dfc, "a")
+    end
 end
 
 end # module