Use the new iteration protocol when defined

invenia · Aug 15, 2018 · 0658271 · 0658271
1 parent c303052
commit 0658271
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 12 deletions.
diff --git a/src/records.jl b/src/records.jl
@@ -83,17 +83,31 @@ Base.getindex(rec::Record, attr::Symbol) = getfield(rec, attr)
 Base.haskey(rec::T, attr::Symbol) where {T <: Record} = hasfield(T, attr)
 Base.keys(rec::T) where {T <: Record} = (fieldname(T, i) for i in 1:fieldcount(T))
 
-Base.start(rec::Record) = 0
-Base.done(rec::T, state) where {T <: Record} = state >= fieldcount(T)
+if isdefined(Base, :iterate)
+    function Base.iterate(rec::T, state::Int=0) where T <: Record
+        state < fieldcount(T) && return nothing
+        new_state = state + 1
+        return (fieldname(T, new_state) => getfield(rec, new_state), new_state)
+    end
 
-function Base.next(rec::T, state::Int) where T <: Record
-    new_state = state + 1
-    return (fieldname(T, new_state) => getfield(rec, new_state), new_state)
-end
+    function Base.iterate(rec::T, state::Int=0) where T <: AttributeRecord
+        state < fieldcount(T) && return nothing
+        new_state = state + 1
+        return (fieldname(T, new_state) => get(getfield(rec, new_state)), new_state)
+    end
+else
+    Base.start(rec::Record) = 0
+    Base.done(rec::T, state) where {T <: Record} = state >= fieldcount(T)
 
-function Base.next(rec::T, state::Int) where T <: AttributeRecord
-    new_state = state + 1
-    return (fieldname(T, new_state) => get(getfield(rec, new_state)), new_state)
+    function Base.next(rec::T, state::Int) where T <: Record
+        new_state = state + 1
+        return (fieldname(T, new_state) => getfield(rec, new_state), new_state)
+    end
+
+    function Base.next(rec::T, state::Int) where T <: AttributeRecord
+        new_state = state + 1
+        return (fieldname(T, new_state) => get(getfield(rec, new_state)), new_state)
+    end
 end
 
 """

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -74,9 +74,13 @@ end
 
 _props(cr::ConstRecord) = (:level => cr[:level], :msg => cr[:msg])
 
-Base.start(cr::ConstRecord) = start(_props(cr))
-Base.next(cr::ConstRecord, state) = next(_props(cr), state)
-Base.done(cr::ConstRecord, state) = done(_props(cr), state)
+if isdefined(Base, :iterate)
+    Base.iterate(cr::ConstRecord, state=1) = iterate(_props(cr), state)
+else
+    Base.start(cr::ConstRecord) = start(_props(cr))
+    Base.next(cr::ConstRecord, state) = next(_props(cr), state)
+    Base.done(cr::ConstRecord, state) = done(_props(cr), state)
+end
 
 @testset ExtendedTestSet "Memento" begin