Fix eltype of flatten of tuple with non-2 length (#56802)

In 4c076c8, eltype of flatten of tuple was improved by computing a refined eltype at compile time. However, this implementation only worked for length-2 tuples, and errored for all others. Generalize this to all tuples. Closes #56783 --------- Co-authored-by: Neven Sajko <[email protected]>
JuliaLang · Jan 2, 2025 · 8fc55be · 8fc55be
1 parent d7ff57b
commit 8fc55be
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/base/iterators.jl b/base/iterators.jl
@@ -15,7 +15,8 @@ using .Base:
     AbstractRange, AbstractUnitRange, UnitRange, LinearIndices, TupleOrBottom,
     (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing,
     any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex,
-    tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString
+    tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString,
+    afoldl
 using Core: @doc
 
 using .Base:
@@ -1202,8 +1203,13 @@ julia> [(x,y) for x in 0:1 for y in 'a':'c']  # collects generators involving It
 flatten(itr) = Flatten(itr)
 
 eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I))
-eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} = promote_typejoin(map(eltype, fieldtypes(I))...)
-eltype(::Type{Flatten{Tuple{}}}) = eltype(Tuple{})
+
+# For tuples, we statically know the element type of each index, so we can compute
+# this at compile time.
+function eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}}
+    afoldl((T, i) -> promote_typejoin(T, eltype(i)), Union{}, fieldtypes(I)...)
+end
+
 IteratorEltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, IteratorEltype(I))
 IteratorEltype(::Type{Flatten{Tuple{}}}) = IteratorEltype(Tuple{})
 _flatteneltype(I, ::HasEltype) = IteratorEltype(eltype(I))

diff --git a/test/iterators.jl b/test/iterators.jl
@@ -515,6 +515,14 @@ end
 @test eltype(flatten(UnitRange{Int8}[1:2, 3:4])) == Int8
 @test eltype(flatten(([1, 2], [3.0, 4.0]))) == Real
 @test eltype(flatten((a = [1, 2], b = Int8[3, 4]))) == Signed
+@test eltype(flatten((Int[], Nothing[], Int[]))) == Union{Int, Nothing}
+@test eltype(flatten((String[],))) == String
+@test eltype(flatten((Int[], UInt[], Int8[],))) == Integer
+@test eltype(flatten((; a = Int[], b = Nothing[], c = Int[]))) == Union{Int, Nothing}
+@test eltype(flatten((; a = String[],))) == String
+@test eltype(flatten((; a = Int[], b = UInt[], c = Int8[],))) == Integer
+@test eltype(flatten(())) == Union{}
+@test eltype(flatten((;))) == Union{}
 @test length(flatten(zip(1:3, 4:6))) == 6
 @test length(flatten(1:6)) == 6
 @test collect(flatten(Any[])) == Any[]