Add getindex/setindex! methods for Colon

Needed since JuliaLang/julia#10331
JuliaData · Jun 24, 2015 · 4e3ebe3 · 4e3ebe3
1 parent fd88acf
commit 4e3ebe3
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 18 deletions.
diff --git a/src/dataframe/dataframe.jl b/src/dataframe/dataframe.jl
@@ -240,6 +240,9 @@ function Base.getindex{T <: ColumnIndex}(df::DataFrame, col_inds::AbstractVector
     return DataFrame(new_columns, Index(_names(df)[selected_columns]))
 end
 
+# df[:] => (Sub)?DataFrame
+Base.getindex(df::DataFrame, col_inds::Colon) = copy(df)
+
 # df[SingleRowIndex, SingleColumnIndex] => Scalar
 function Base.getindex(df::DataFrame, row_ind::Real, col_ind::ColumnIndex)
     selected_column = index(df)[col_ind]
@@ -266,6 +269,22 @@ function Base.getindex{R <: Real, T <: ColumnIndex}(df::DataFrame, row_inds::Abs
     return DataFrame(new_columns, Index(_names(df)[selected_columns]))
 end
 
+# df[:, SingleColumnIndex] => (Sub)?AbstractVector
+# df[:, MultiColumnIndex] => (Sub)?DataFrame
+Base.getindex{T<:ColumnIndex}(df::DataFrame, row_inds::Colon, col_inds::Union(T, AbstractVector{T})) = df[col_inds]
+
+# df[SingleRowIndex, :] => (Sub)?DataFrame
+Base.getindex(df::DataFrame, row_ind::Real, col_inds::Colon) = df[[row_ind], col_inds]
+
+# df[MultiRowIndex, :] => (Sub)?DataFrame
+function Base.getindex{R<:Real}(df::DataFrame, row_inds::AbstractVector{R}, col_inds::Colon)
+    new_columns = Any[dv[row_inds] for dv in df.columns]
+    return DataFrame(new_columns, copy(index(df)))
+end
+
+# df[:, :] => (Sub)?DataFrame
+Base.getindex(df::DataFrame, ::Colon, ::Colon) = copy(df)
+
 ##############################################################################
 ##
 ## setindex!()
@@ -403,6 +422,9 @@ function Base.setindex!{T <: ColumnIndex}(df::DataFrame,
     return df
 end
 
+# df[:] = AbstractVector or Single Item
+Base.setindex!(df::DataFrame, v, ::Colon) = (df[1:size(df, 2)] = v; df)
+
 # df[SingleRowIndex, SingleColumnIndex] = Single Item
 function Base.setindex!(df::DataFrame,
                 v::Any,
@@ -562,6 +584,28 @@ function Base.setindex!{R <: Real, T <: ColumnIndex}(df::DataFrame,
     return df
 end
 
+# df[:] = DataFrame, df[:, :] = DataFrame
+function Base.setindex!(df::DataFrame,
+                                  new_df::DataFrame,
+                                  row_inds::Colon,
+                                  col_inds::Colon=Colon())
+    df.columns = copy(new_df.columns)
+    df.colindex = copy(new_df.colindex)
+    df
+end
+
+# df[:, :] = ...
+Base.setindex!(df::DataFrame, v, ::Colon, ::Colon) =
+    (df[1:size(df, 1), 1:size(df, 2)] = v; df)
+
+# df[Any, :] = ...
+Base.setindex!(df::DataFrame, v, row_inds, ::Colon) =
+    (df[row_inds, 1:size(df, 2)] = v; df)
+
+# df[:, Any] = ...
+Base.setindex!(df::DataFrame, v, ::Colon, col_inds) =
+    (df[col_inds] = v; df)
+
 # Special deletion assignment
 Base.setindex!(df::DataFrame, x::Nothing, col_ind::Int) = delete!(df, col_ind)
 

diff --git a/test/indexing.jl b/test/indexing.jl
@@ -84,31 +84,33 @@ module TestIndexing
 
     df = DataFrame(A = 1:2, B = 3:4)
 
-    row_indices = {1,
-                   1.0,
-                   1:2,
-                   [1, 2],
-                   [1.0, 2.0],
-                   [true, true],
-                   trues(2),
-                   @data([1, 2]),
-                   @data([1.0, 2.0]),
-                   @data([true, true])}
-
-    column_indices = {1,
+    row_indices = Any[1,
                       1.0,
-                      "A",
                       1:2,
                       [1, 2],
                       [1.0, 2.0],
-                      [true, false],
+                      [true, true],
                       trues(2),
-                      ["A", "B"],
                       @data([1, 2]),
                       @data([1.0, 2.0]),
-                      @data([true, false]),
-                      @data(["A", "B"]),
-                      :(names(_DF) .== "B")}
+                      @data([true, true]),
+                      Colon()]
+
+    column_indices = Any[1,
+                         1.0,
+                         "A",
+                         1:2,
+                         [1, 2],
+                         [1.0, 2.0],
+                         [true, false],
+                         trues(2),
+                         ["A", "B"],
+                         @data([1, 2]),
+                         @data([1.0, 2.0]),
+                         @data([true, false]),
+                         @data(["A", "B"]),
+                         :(names(_DF) .== "B"),
+                         Colon()]
 
     #
     # getindex()