diff --git a/NEWS.md b/NEWS.md index f9fb0835afffc..ffbc96a4d3d35 100644 --- a/NEWS.md +++ b/NEWS.md @@ -51,6 +51,7 @@ Standard library changes constructing the range. ([#40382]) * TCP socket objects now expose `closewrite` functionality and support half-open mode usage ([#40783]). * Intersect returns a result with the eltype of the type-promoted eltypes of the two inputs ([#41769]). +* `Iterators.countfrom` now accepts any type that defines `+`. ([#37747]) #### InteractiveUtils * A new macro `@time_imports` for reporting any time spent importing packages and their dependencies ([#41612]) diff --git a/base/iterators.jl b/base/iterators.jl index c0a68a36d836e..5116bd450021f 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -592,8 +592,8 @@ IteratorSize(::Type{<:Rest{I}}) where {I} = rest_iteratorsize(IteratorSize(I)) # Count -- infinite counting -struct Count{S<:Number} - start::S +struct Count{T,S} + start::T step::S end @@ -613,11 +613,13 @@ julia> for v in Iterators.countfrom(5, 2) 9 ``` """ -countfrom(start::Number, step::Number) = Count(promote(start, step)...) -countfrom(start::Number) = Count(start, oneunit(start)) -countfrom() = Count(1, 1) +countfrom(start::T, step::S) where {T,S} = Count{typeof(start+step),S}(start, step) +countfrom(start::Number, step::Number) = Count(promote(start, step)...) +countfrom(start) = Count(start, oneunit(start)) +countfrom() = Count(1, 1) -eltype(::Type{Count{S}}) where {S} = S + +eltype(::Type{Count{T,S}}) where {T,S} = T iterate(it::Count, state=it.start) = (state, state + it.step) diff --git a/test/iterators.jl b/test/iterators.jl index c7d00c4e7e2e8..010697f5ae4e2 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -3,6 +3,7 @@ using Base.Iterators using Random using Base: IdentityUnitRange +using Dates: Date, Day @test Base.IteratorSize(Any) isa Base.SizeUnknown @@ -123,7 +124,7 @@ end # countfrom # --------- -let i = 0, k = 1 +let i = 0, k = 1, l = 0 for j = countfrom(0, 2) @test j == i*2 i += 1 @@ -134,6 +135,15 @@ let i = 0, k = 1 k += 1 k <= 10 || break end + # test that `start` promotes to `typeof(start+step)` + for j = countfrom(Int[0, 0], Float64[1.0, 2.0]) + @test j isa Vector{Float64} + @test j == l*[1, 2] + l += 1 + l <= 10 || break + end + # test with `start` and `step` having different types + @test collect(take(countfrom(Date(2020,12,25), Day(1)), 12)) == range(Date(2020,12,25), step=Day(1), length=12) end # take