Merge pull request #80 from Tokazama/master

Add known_length method for generators

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.13.6"
+version = "2.13.7"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/ArrayInterface.jl

@@ -36,6 +36,7 @@ known_length(x) = known_length(typeof(x))
 known_length(::Type{<:NamedTuple{L}}) where {L} = length(L)
 known_length(::Type{T}) where {T<:Slice} = known_length(parent_type(T))
 known_length(::Type{<:Tuple{Vararg{Any,N}}}) where {N} = N
+known_length(::Type{T}) where {Itr,T<:Base.Generator{Itr}} = known_length(Itr)
 known_length(::Type{<:Number}) = 1
 function known_length(::Type{T}) where {T}
     if parent_type(T) <: T
@@ -50,6 +51,7 @@ end
 _known_length(x::Tuple{Vararg{<:Union{Int,Nothing}}}) = nothing
 _known_length(x::Tuple{Vararg{Int}}) = prod(x)
 
+
 """
     can_change_size(::Type{T}) -> Bool
 

src/indexing.jl

@@ -218,7 +218,7 @@ to_indices(A, axs::Tuple{}, args::Tuple{}) = ()
 
 
 _multi_check_index(axs::Tuple, arg) = _multi_check_index(axs, axes(arg))
-_multi_check_index(axs::Tuple, arg::LogicalIndex) = prod(map(length, axs)) == length(arg)
+_multi_check_index(axs::Tuple, arg::LogicalIndex) = axs == axes(arg.mask)
 function _multi_check_index(axs::Tuple, arg::AbstractArray{T}) where {T<:CartesianIndex}
     b = true
     for i in arg

test/runtests.jl

@@ -411,6 +411,9 @@ end
     x = view(SArray{Tuple{3,3,3}}(ones(3,3,3)), :, SOneTo(2), 2)
     @test @inferred(ArrayInterface.known_length(x)) == 6
     @test @inferred(ArrayInterface.known_length(x')) == 6
+
+    itr = StaticInt(1):StaticInt(10)
+    @inferred(ArrayInterface.known_length((i for i in itr))) == 10
 end
 
 @testset "indices" begin