Define shape of some truncated zips (#29927)

NEWS.md

@@ -104,6 +104,7 @@ New library features
 
 Standard library changes
 ------------------------
+* A 1-d `Zip` iterator (where `Base.IteratorSize` is `Base.HasShape{1}()`) with defined length of `n` has now also size of `(n,)` (instead of throwing an error with truncated iterators) ([#29927]).
 * The `@timed` macro now returns a `NamedTuple` ([#34149])
 * New `supertypes(T)` function returns a tuple of all supertypes of `T` ([#34419]).
 * Sorting-related functions such as `sort` that take the keyword arguments `lt`, `rev`, `order`
@@ -157,4 +158,5 @@ Tooling Improvements
 ---------------------
 
 
+
 <!--- generated by NEWS-update.jl: -->

base/iterators.jl

@@ -313,10 +313,11 @@ function _zip_min_length(is)
     end
 end
 _zip_min_length(is::Tuple{}) = nothing
-size(z::Zip) = _promote_shape(map(size, z.is)...)
-axes(z::Zip) = _promote_shape(map(axes, z.is)...)
-_promote_shape(a, b...) = promote_shape(a, _promote_shape(b...))
-_promote_shape(a) = a
+size(z::Zip) = mapreduce(size, _zip_promote_shape, z.is)
+axes(z::Zip) = mapreduce(axes, _zip_promote_shape, z.is)
+_zip_promote_shape((a,)::Tuple{OneTo}, (b,)::Tuple{OneTo}) = (intersect(a, b),)
+_zip_promote_shape((m,)::Tuple{Integer},(n,)::Tuple{Integer}) = (min(m,n),)
+_zip_promote_shape(a, b) = promote_shape(a, b)
 eltype(::Type{Zip{Is}}) where {Is<:Tuple} = _zip_eltype(Is)
 _zip_eltype(::Type{Is}) where {Is<:Tuple} =
     tuple_type_cons(eltype(tuple_type_head(Is)), _zip_eltype(tuple_type_tail(Is)))

test/iterators.jl

@@ -16,10 +16,44 @@ let z = zip(1:2)
     @test eltype(z) == Tuple{Int}
 end
 
-let z = zip(1:2, 3:4)
+for z in (zip(1:2, 3:4), zip(1:2, 3:5))
+    @test collect(z) == [(1,3), (2,4)]
+    @test eltype(z) == Tuple{Int,Int}
     @test size(z) == (2,)
+    @test axes(z) == (Base.OneTo(2),)
+    @test length(z) == 2
+end
+
+let z = zip(1:2, Iterators.countfrom(3))
     @test collect(z) == [(1,3), (2,4)]
     @test eltype(z) == Tuple{Int,Int}
+    @test_throws MethodError size(z) # by convention, the zip of a finite and
+                         # an infinite iterator has only `length`
+    @test_throws MethodError axes(z)
+    @test length(z) == 2
+end
+
+let z = zip([i*j for i in 1:3, j in -1:2:1], 1:6)
+    @test collect(z) == [(-1, 1)
+                         (-2, 2)
+                         (-3, 3)
+                         (1, 4)
+                         (2, 5)
+                         (3, 6) ]
+    @test eltype(z) == Tuple{Int,Int}
+    @test_throws DimensionMismatch size(z)
+    @test_throws DimensionMismatch axes(z)
+    @test length(z) == 6
+end
+
+let z = zip([i*j for i in 1:3, j in -1:2:1], [i*j for i in 1:3, j in -1:2:1])
+    @test collect(z) == [(-1, -1) (1, 1)
+                        (-2, -2) (2, 2)
+                        (-3, -3) (3, 3)]
+    @test eltype(z) == Tuple{Int,Int}
+    @test size(z) == (3, 2)
+    @test axes(z) == (Base.OneTo(3), Base.OneTo(2))
+    @test length(z) == 6
 end
 
 let z = zip(1:2, 3:4, 5:6)