diff --git a/src/abstractdataarray.jl b/src/abstractdataarray.jl index 6381e61..3629ef0 100644 --- a/src/abstractdataarray.jl +++ b/src/abstractdataarray.jl @@ -157,8 +157,13 @@ type EachReplaceNA{S, T} replacement::T end function each_replacena(da::AbstractDataArray, replacement::Any) + if isa(replacement, Function) + EachReplaceNAWithFunctionResult(da, replacement) + else EachReplaceNA(da, convert(eltype(da), replacement)) + end end + function each_replacena(replacement::Any) x -> each_replacena(x, replacement) end @@ -167,4 +172,16 @@ Base.done(itr::EachReplaceNA, ind::Integer) = ind > length(itr.da) function Base.next(itr::EachReplaceNA, ind::Integer) item = isna(itr.da, ind) ? itr.replacement : itr.da[ind] (item, ind + 1) -end \ No newline at end of file +end + +type EachReplaceNAWithFunctionResult{T} + da::AbstractDataArray{T} + f::Function +end + +Base.start(itr::EachReplaceNAWithFunctionResult) = 1 +Base.done(itr::EachReplaceNAWithFunctionResult, ind::Int) = ind > length(itr.da) +function Base.next(itr::EachReplaceNAWithFunctionResult, ind::Integer) + item = isna(itr.da, ind) ? itr.f(itr.da) : itr.da[ind] + (item, ind + 1) +end diff --git a/test/nas.jl b/test/nas.jl index c4df9e5..7c0107c 100644 --- a/test/nas.jl +++ b/test/nas.jl @@ -63,4 +63,9 @@ module TestNAs @test_throws NAException for v in each_failna(dv); end @test collect(each_dropna(dv)) == a @test collect(each_replacena(dv, 4)) == [4, 4, a..., 4] + + dv = DataArray(Array(1:6), push!(fill(false, 3),fill(true, 3)...)) + a = dropna(dv) + f = x -> round(Int, mean(dropna(x))) + @test collect(each_replacena(dv, f)) == [1, 2, 3, 2, 2, 2] end