Supporting indexing EachReplaceMissing and EachFailMissing (#126)

Project.toml

@@ -1,6 +1,6 @@
 name = "Missings"
 uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
-version = "0.4.4"
+version = "0.4.5"
 
 [deps]
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"

src/Missings.jl

@@ -53,10 +53,15 @@ disallowmissing(x::AbstractArray{T}) where {T} = convert(AbstractArray{nonmissin
 """
     Missings.replace(itr, replacement)
 
-Return an iterator wrapping iterable `itr` which replaces [`missing`](@ref) values with
-`replacement`. When applicable, the size of `itr` is preserved.
+Return an iterator over the elements in `itr` which replaces [`missing`](@ref) values with
+`replacement`.
+
+When applicable, the size of `itr` is preserved.
+The returned object can be indexed using indices of `itr` if the latter is indexable,
+and `eachindex` and `keys` return the indices of `itr`.
+
 If the type of `replacement` differs from the element type of `itr`,
-it will be converted.
+it will be converted to it.
 
 See also: [`skipmissing`](@ref), [`Missings.fail`](@ref)
 
@@ -86,7 +91,10 @@ Base.IteratorEltype(::Type{<:EachReplaceMissing{T}}) where {T} =
     Base.IteratorEltype(T)
 Base.length(itr::EachReplaceMissing) = length(itr.x)
 Base.size(itr::EachReplaceMissing) = size(itr.x)
+Base.axes(itr::EachReplaceMissing) = axes(itr.x)
 Base.eltype(itr::EachReplaceMissing) = nonmissingtype(eltype(itr.x))
+Base.eachindex(itr::EachReplaceMissing) = eachindex(itr.x)
+Base.keys(itr::EachReplaceMissing) = keys(itr.x)
 
 @inline function Base.iterate(itr::EachReplaceMissing)
     st = iterate(itr.x)
@@ -102,11 +110,19 @@ end
     return (v isa Missing ? itr.replacement : v, s)
 end
 
+Base.@propagate_inbounds function Base.getindex(itr::EachReplaceMissing, I...)
+    v = itr.x[I...]
+    return v === missing ? itr.replacement : v
+end
+
 """
     Missings.fail(itr)
 
-Return an iterator wrapping iterable `itr` which will throw a [`MissingException`](@ref)
+Return an iterator over the elements in `itr` which will throw a [`MissingException`](@ref)
 if a [`missing`](@ref) value is found.
+The returned object can be indexed using indices of `itr` if the latter is indexable.
+Indices corresponding to missing values are not valid: even though they are not skipped
+by `keys` and eachindex, a `MissingException` is thrown when trying to use them.
 
 Use [`collect`](@ref) to obtain an `Array` containing the resulting values.
 If `itr` is an array, the resulting array will have the same dimensions.
@@ -135,7 +151,10 @@ Base.IteratorEltype(::Type{EachFailMissing{T}}) where {T} =
     Base.IteratorEltype(T)
 Base.length(itr::EachFailMissing) = length(itr.x)
 Base.size(itr::EachFailMissing) = size(itr.x)
+Base.axes(itr::EachFailMissing) = axes(itr.x)
 Base.eltype(itr::EachFailMissing) = nonmissingtype(eltype(itr.x))
+Base.eachindex(itr::EachFailMissing) = eachindex(itr.x)
+Base.keys(itr::EachFailMissing) = keys(itr.x)
 
 @inline function Base.iterate(itr::EachFailMissing)
     st = iterate(itr.x)
@@ -153,6 +172,12 @@ end
     return (v::eltype(itr), s)
 end
 
+Base.@propagate_inbounds function Base.getindex(itr::EachFailMissing, I...)
+    v = itr.x[I...]
+    v === missing && throw(MissingException("the value at index $I is missing"))
+    return v
+end
+
 const levels = DataAPI.levels
 
 struct PassMissing{F} <: Function

test/runtests.jl

@@ -9,23 +9,37 @@ struct CubeRooter end
     @test eltype(x) === Int
     @test length(x) == 4
     @test size(x) == (4,)
-    @test collect(x) == collect(1:4)
+    @test axes(x) == (1:4,)
+    @test keys(x) == eachindex(x) == 1:4
+    @test collect(x) == [x[i] for i in keys(x)] == 1:4
     @test collect(x) isa Vector{Int}
     x = Missings.replace([1, 2, missing, 4], 3.0)
     @test eltype(x) === Int
     @test length(x) == 4
     @test size(x) == (4,)
-    @test collect(x) == collect(1:4)
+    @test axes(x) == (1:4,)
+    @test keys(x) == eachindex(x) == 1:4
+    @test collect(x) == [x[i] for i in keys(x)] == 1:4
     @test collect(x) isa Vector{Int}
     x = Missings.replace([1 2; missing 4], 3)
     @test eltype(x) === Int
     @test length(x) == 4
     @test size(x) == (2, 2)
-    @test collect(x) == [1 2; 3 4]
+    @test axes(x) == (1:2, 1:2)
+    @test keys(x) == CartesianIndices((1:2, 1:2))
+    @test eachindex(x) == 1:4
+    @test collect(x) == [x[i] for i in keys(x)] == [1 2; 3 4]
     @test collect(x) isa Matrix{Int}
     x = Missings.replace((v for v in [missing, 1, missing, 2, 4]), 0)
     @test length(x) == 5
     @test size(x) == (5,)
+    if VERSION >= v"1.7.0-DEV"
+        @test eachindex(x) == keys(x) == 1:5
+    else
+        @test_throws MethodError keys(x)
+        @test_throws MethodError eachindex(x)
+    end
+    @test_throws MethodError x[1]
     @test eltype(x) === Any
     @test collect(x) == [0, 1, 0, 2, 4]
     @test collect(x) isa Vector{Int}
@@ -34,19 +48,32 @@ struct CubeRooter end
     @test eltype(x) === Int
     @test length(x) == 4
     @test size(x) == (4,)
-    @test collect(x) == [1, 2, 3, 4]
+    @test axes(x) == (1:4,)
+    @test keys(x) == eachindex(x) == 1:4
+    @test collect(x) == [x[i] for i in keys(x)] == [1, 2, 3, 4]
     @test collect(x) isa Vector{Int}
     x = Missings.fail([1 2; 3 4])
     @test eltype(x) === Int
     @test length(x) == 4
     @test size(x) == (2, 2)
-    @test collect(x) == [1 2; 3 4]
+    @test axes(x) == (1:2, 1:2)
+    @test keys(x) == CartesianIndices((1:2, 1:2))
+    @test eachindex(x) == 1:4
+    @test collect(x) == [x[i] for i in keys(x)] == [1 2; 3 4]
     @test collect(x) isa Matrix{Int}
     @test_throws MissingException collect(Missings.fail([1, 2, missing, 4]))
     x = Missings.fail(v for v in [1, 2, 4])
     @test eltype(x) === Any
     @test length(x) == 3
     @test size(x) == (3,)
+    @test axes(x) == (1:3,)
+    if VERSION >= v"1.7.0-DEV"
+        @test eachindex(x) == keys(x) == 1:3
+    else
+        @test_throws MethodError keys(x)
+        @test_throws MethodError eachindex(x)
+    end
+    @test_throws MethodError x[1]
     @test collect(x) == [1, 2, 4]
     @test collect(x) isa Vector{Int}