Add getindex/setindex! methods for Colon

Needed since JuliaLang/julia#10331
JuliaData · Jun 24, 2015 · c793c5b · c793c5b
1 parent fd88acf
commit c793c5b
Showing 1 changed file with 44 additions and 0 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)