Add setindex!! support for DataFrames

ext/BangBangDataFramesExt.jl

@@ -14,6 +14,22 @@ end
     BangBang.append!!(df::DataFrames.DataFrame, source) = df_append!!(df, source)
     BangBang.copyappendable(df::DataFrames.AbstractDataFrame) = copy(df)
 
+    function BangBang.setindex!!(
+        df::DataFrames.DataFrame,
+        v,
+        row::Integer,
+        col::Union{Integer,Symbol,AbstractString},
+    )
+        col_idx = col isa Integer ? Int(col) : DataFrames.columnindex(df, col)
+        columns = getfield(df, :columns)
+        old_col = columns[col_idx]
+        new_col = BangBang.setindex!!(old_col, v, row)
+        if new_col !== old_col
+            columns[col_idx] = new_col
+        end
+        return df
+    end
+
     _getvalue(x, pos, name) = getproperty(x, name)
     _getvalue(x::AbstractVector, pos, name) = x[pos]
     _getvalue(x::Tuple, pos, name) = x[pos]

test/test_dataframes.jl

@@ -1,6 +1,6 @@
 module TestDataFrames
 
-using BangBang: append!!, push!!
+using BangBang: append!!, push!!, setindex!!
 using CategoricalArrays: CategoricalArray
 using DataFrames: DataFrame
 using InitialValues: InitialValue
@@ -84,4 +84,32 @@ end
     end
 end
 
+@testset "setindex!!" begin
+    @testset "basic" begin
+        df = DataFrame("a" => [0.0])
+        @test setindex!!(df, 1.0, 1, "a") === df
+        @test df[1, "a"] === 1.0
+    end
+
+    @testset "Symbol column" begin
+        df = DataFrame(a = [0.0])
+        @test setindex!!(df, 2.0, 1, :a) === df
+        @test df[1, :a] === 2.0
+    end
+
+    @testset "Integer column" begin
+        df = DataFrame(a = [0.0])
+        @test setindex!!(df, 3.0, 1, 1) === df
+        @test df[1, 1] === 3.0
+    end
+
+    @testset "type widening" begin
+        df = DataFrame(a = [1, 2, 3])
+        result = setindex!!(df, 1.5, 2, :a)
+        @test result === df
+        @test df[2, :a] === 1.5
+        @test eltype(df.a) === Float64
+    end
+end
+
 end  # module