Proper definition of eltype for custom iterator

I was looking at the Julia manual and read about custom iterators. The gist: having defined my own type, once I extend Base.iterate + Base.length + Base.eltype to work with my own type, I can do cool stuff.

Here’s the example from the manual, slightly modified:

struct Squares
    count::Int
    eltype::DataType
    function Squares(c::Int,e::DataType)
        e <: Real ? new(c,e) : throw(error("only real DataTypes allowed"))
    end
end

Base.iterate(s::Squares, state=1) = state > s.count ? nothing : ( s.eltype(state*state), state+1 )
Base.length(s::Squares) = s.count
Base.eltype(s::Squares) = s.eltype # is this the best way to do this?

This let’s me do the following out-of-the-box:

julia> for i in Squares(4,UInt8) println(i) end
1
4
9
16

julia> collect(Squares(3,Rational{Int64}))
3-element Array{Rational{Int64},1}:
 1//1
 4//1
 9//1

julia> 9 in Squares(3,Float32)
true

This is cool stuff!

My question is: is Base.eltype(s::Squares) = s.eltype the best way to do this? I know the answer is no, quoting from the manual:

The definition eltype(x) = eltype(typeof(x)) is provided for convenience so that instances can be passed instead of types. However the form that accepts a type argument should be defined for new types.

So what’s the correct syntax for Base.eltype(::Type{Squares}) = Squares.eltype (which is clearly incorrect!)?

cheers!

I would write it as:

struct Squares{T <: Real}
    count::Int
end

Base.iterate(s::Squares{T}, state=1) where {T} = state > s.count ? nothing : (T(state^2), state+1)
Base.length(s::Squares) = s.count
Base.eltype(::Type{Squares{T}}) where {T} = T

And the code you executed would be written as:

for i in Squares{UInt8}(4) println(i) end

collect(Squares{Rational{Int64}}(3))

9 in Squares{Float32}(3)
2 Likes

Cool! This looks good to me. Thansk!