Skip to content

Commit

Permalink
fix getindex and iterate
Browse files Browse the repository at this point in the history
  • Loading branch information
ogauthe committed Nov 7, 2024
1 parent 0d51b19 commit f6911ee
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 0 deletions.
9 changes: 9 additions & 0 deletions NDTensors/src/lib/GradedAxes/src/labelledunitrangedual.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,15 @@ end

# fix ambiguities
Base.getindex(a::LabelledUnitRangeDual, i::Integer) = dual(nondual(a)[i])
function Base.getindex(a::LabelledUnitRangeDual, indices::AbstractUnitRange{<:Integer})
return dual(nondual(a)[indices])
end

function Base.iterate(a::LabelledUnitRangeDual, i)
i == last(a) && return nothing
next = convert(eltype(a), labelled(i + step(a), label(a)))
return (next, next)
end

function Base.show(io::IO, ::MIME"text/plain", a::LabelledUnitRangeDual)
println(io, typeof(a))
Expand Down
17 changes: 17 additions & 0 deletions NDTensors/src/lib/GradedAxes/test/test_dual.jl
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,11 @@ end
@test space_isequal(la, la)
@test label_type(la) == U1

@test iterate(la) == (1, 1)
@test iterate(la) == (1, 1)
@test iterate(la, 1) == (2, 2)
@test isnothing(iterate(la, 2))

lad = dual(la)
@test lad isa LabelledUnitRangeDual
@test label(lad) == U1(-1)
Expand All @@ -87,10 +92,22 @@ end
@test dual(lad) === la
@test label_type(lad) == U1

@test iterate(lad) == (1, 1)
@test iterate(lad) == (1, 1)
@test iterate(lad, 1) == (2, 2)
@test isnothing(iterate(lad, 2))

lad2 = lad[1:1]
@test lad2 isa LabelledUnitRangeDual
@test label(lad2) == U1(-1)
@test unlabel(lad2) == 1:1

laf = flip(la)
@test laf isa LabelledUnitRangeDual
@test label(laf) == U1(1)
@test unlabel(laf) == 1:2
@test labelled_isequal(la, laf)
@test !space_isequal(la, laf)

ladf = flip(dual(la))
@test ladf isa LabelledUnitRange
Expand Down

0 comments on commit f6911ee

Please sign in to comment.