add Type to names as cols

NEWS.md

@@ -55,6 +55,8 @@
   ([#2373](https://github.com/JuliaData/DataFrames.jl/pull/2373))
 * add `columnindex` for `DataFrameRow`
   ([#2380](https://github.com/JuliaData/DataFrames.jl/pull/2380))
+* `names` now accepts `Type` as a column selector
+  ([#2400](https://github.com/JuliaData/DataFrames.jl/pull/2400))
 * `select`, `select!`, `transform`, `transform!` and `combine` now allow `renamecols`
   keyword argument that makes it possible to avoid adding transformation function name
   as a suffix in automatically generated column names

src/abstractdataframe/abstractdataframe.jl

@@ -68,7 +68,8 @@ Return a freshly allocated `Vector{String}` of names of columns contained in `df
 
 If `cols` is passed then restrict returned column names to those matching the
 selector (this is useful in particular with regular expressions, `Not`, and `Between`).
-`cols` can be any column selector ($COLUMNINDEX_STR; $MULTICOLUMNINDEX_STR).
+`cols` can be any column selector ($COLUMNINDEX_STR; $MULTICOLUMNINDEX_STR)
+or a `Type`, in which case columns whose `eltype` is a subtype of `cols` are returned.
 
 See also [`propertynames`](@ref) which returns a `Vector{Symbol}`.
 """
@@ -78,9 +79,12 @@ function Base.names(df::AbstractDataFrame, cols)
     nms = _names(index(df))
     idx = index(df)[cols]
     idxs = idx isa Int ? (idx:idx) : idx
-    return [string(nms[i]) for i in idxs]
+    return [String(nms[i]) for i in idxs]
 end
 
+Base.names(df::AbstractDataFrame, T::Type) =
+    [String(n) for (n, c) in pairs(eachcol(df)) if eltype(c) <: T]
+
 # _names returns Vector{Symbol} without copying
 _names(df::AbstractDataFrame) = _names(index(df))
 

src/deprecated.jl

@@ -49,15 +49,14 @@ function CategoricalArrays.categorical(df::AbstractDataFrame,
     if cols === nothing
         cols = Union{AbstractString, Missing}
         Base.depwarn("`categorical(df)` is deprecated. " *
-                     "Use `cols = names(df)[map(c -> eltype(c) <: $cols, eachcol(df))]; transform(df, cols .=> $categoricalstr, renamecols=false)` instead.",
+                     "Use `transform(df, names(df, $cols) .=> $categoricalstr, renamecols=false)` instead.",
                      :categorical)
     else
         Base.depwarn("`categorical(df, T)` is deprecated. " *
-                     "Use `cols = names(df)[map(c -> eltype(c) <: T, eachcol(df))]; transform(df, cols .=> $categoricalstr, renamecols=false)` instead.",
+                     "Use transform(df, names(df, T) .=> $categoricalstr, renamecols=false)` instead.",
                      :categorical)
     end
-    colsstr = names(df)[map(c -> eltype(c) <: cols, eachcol(df))]
-    return transform(df, colsstr .=> (x -> categorical(x, compress=compress)), renamecols=false)
+    return transform(df, names(df, cols) .=> (x -> categorical(x, compress=compress)), renamecols=false)
 end
 
 function categorical!(df::DataFrame, cols::Union{ColumnIndex, MultiColumnIndex};
@@ -97,13 +96,12 @@ function categorical!(df::DataFrame, cols::Union{Type, Nothing}=nothing;
     if cols === nothing
         cols = Union{AbstractString, Missing}
         Base.depwarn("`categorical!(df)` is deprecated. " *
-                     "Use `cols = names(df)[map(c -> eltype(c) <: $cols, eachcol(df))]; transform!(df, cols .=> $categoricalstr, renamecols=false)` instead.",
+                     "Use `transform!(df, names(df, $cols) .=> $categoricalstr, renamecols=false)` instead.",
                      :categorical!)
     else
         Base.depwarn("`categorical!(df, T)` is deprecated. " *
-                     "Use `cols = names(df)[map(c -> eltype(c) <: T, eachcol(df))]; transform!(df, cols .=> $categoricalstr, renamecols=false)` instead.",
+                     "Use `transform!(df, names(df, T) .=> $categoricalstr, renamecols=false)` instead.",
                      :categorical!)
     end
-    colsstr = names(df)[map(c -> eltype(c) <: cols, eachcol(df))]
-    return transform!(df, colsstr .=> (x -> categorical(x, compress=compress)), renamecols=false)
+    return transform!(df, names(df, cols) .=> (x -> categorical(x, compress=compress)), renamecols=false)
 end

test/dataframe.jl

@@ -1881,4 +1881,16 @@ end
     @test_throws ArgumentError push!(df, "a")
 end
 
+@testset "names for Type" begin
+    df = DataFrame(a1 = 1:3, a2 = [1, missing, 3],
+                   b1 = 1.0:3.0, b2 = [1.0, missing, 3.0],
+                   c1 = '1':'3', c2 = ['1', missing, '3'])
+    @test names(df, Int) == ["a1"]
+    @test names(df, Union{Missing, Int}) == ["a1", "a2"]
+    @test names(df, Real) == ["a1", "b1"]
+    @test names(df, Union{Missing, Real}) == ["a1", "a2", "b1", "b2"]
+    @test names(df, Any) == names(df)
+    @test names(df, Union{Char, Float64, Missing}) == ["b1", "b2", "c1", "c2"]
+end
+
 end # module