fix atol in At for Dates (#834)

rafaqz · Nov 2, 2024 · e433038 · e433038
1 parent 9579981
commit e433038
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 16 deletions.
diff --git a/src/Lookups/selector.jl b/src/Lookups/selector.jl
@@ -208,7 +208,7 @@ function at(
     end
 end
 function at(
-    ::Ordered, ::Span, lookup::Lookup{<:Union{Number,Dates.TimeType,AbstractString}}, selval, atol, rtol::Nothing;
+    ::Ordered, ::Span, lookup::Lookup{<:Union{Number,Dates.AbstractTime,AbstractString}}, selval, atol, rtol::Nothing;
     err=_True()
 )
     x = unwrap(selval)
@@ -244,6 +244,8 @@ function at(::Order, ::Span, lookup::Lookup, selval, atol, rtol::Nothing; err=_T
 end
 
 @inline _is_at(x, y, atol) = x == y
+@inline _is_at(x::Dates.AbstractTime, y::Dates.AbstractTime, atol::Dates.Period) = 
+    x >= y - atol && x <= y + atol 
 @inline _is_at(x::Real, y::Real, atol::Real) = abs(x - y) <= atol
 @inline _is_at(x::Real, ys::AbstractArray, atol) = any(y -> _is_at(x, y, atol), ys)
 @inline _is_at(xs::AbstractArray, y::Real, atol) = any(x -> _is_at(x, y, atol), xs)
@@ -349,8 +351,8 @@ end
 _adjust_locus(locus::Center, v, lookup) = v
 _adjust_locus(locus::Start, v, lookup) = v - abs(step(lookup)) / 2
 _adjust_locus(locus::End, v, lookup) = v + abs(step(lookup)) / 2
-_adjust_locus(locus::Start, v::Dates.TimeType, lookup) = v - (v - (v - abs(step(lookup)))) / 2
-_adjust_locus(locus::End, v::Dates.TimeType, lookup) = v + (v + abs(step(lookup)) - v) / 2
+_adjust_locus(locus::Start, v::Dates.AbstractTime, lookup) = v - (v - (v - abs(step(lookup)))) / 2
+_adjust_locus(locus::End, v::Dates.AbstractTime, lookup) = v + (v + abs(step(lookup)) - v) / 2
 _adjust_locus(locus::Start, v::Dates.Date, lookup) = v - (v - (v - abs(step(lookup)))) ÷ 2
 _adjust_locus(locus::End, v::Dates.Date, lookup) = v + (v + abs(step(lookup)) - v) ÷ 2
 
@@ -1172,7 +1174,7 @@ _by(x::AbstractRange) = first(x)
 _by(x::IntervalSets.Interval) = x.left
 _by(x) = x
 
-_in(needle::Dates.TimeType, haystack::Dates.TimeType) = needle == haystack
+_in(needle::Dates.AbstractTime, haystack::Dates.AbstractTime) = needle == haystack
 _in(needle, haystack) = needle in haystack
 _in(needles::Tuple, haystacks::Tuple) = all(map(_in, needles, haystacks))
 _in(needle::Interval, haystack::ClosedInterval) = needle.left in haystack && needle.right in haystack

diff --git a/test/selector.jl b/test/selector.jl
@@ -1177,7 +1177,7 @@ end
         @test view(da, Between(40, 45), At((2:3)u"s")) isa DimArray{Int64,2}
     end
 
-    @testset "with DateTime index" begin
+    @testset "with DateTime lookup" begin
         @testset "Start locus" begin
             timedim = Ti(Sampled(DateTime(2001):Month(1):DateTime(2001, 12); 
                 span=Regular(Month(1)), sampling=Intervals(Start())
@@ -1190,16 +1190,9 @@ end
             @test @inferred da[DateTime(2001, 4, 7) .. DateTime(2001, 8, 30)] == [5, 6, 7]
             @test @inferred da[Date(2001, 4, 7) .. Date(2001, 8, 30)] == [5, 6, 7]
 
-            timedim = Ti(Sampled(Date(2001):Month(1):Date(2001, 12); 
-                span=Regular(Month(1)), sampling=Intervals(Start())
-            ))
-            da = DimArray(1:12, timedim)
-            @test @inferred da[Ti(At(DateTime(2001, 3)))] == 3
-            @test @inferred da[Ti(At(Date(2001, 3)))] == 3
-            @test @inferred da[Near(DateTime(2001, 4, 7))] == 4
-            @test @inferred da[Near(Date(2001, 4, 7))] == 4
-            @test @inferred da[DateTime(2001, 4, 7) .. DateTime(2001, 8, 30)] == [5, 6, 7]
-            @test @inferred da[Date(2001, 4, 7) .. Date(2001, 8, 30)] == [5, 6, 7]
+            @test_throws SelectorError da[Ti(At(Date(2001, 3, 4); atol=Day(2)))]
+            @test @inferred da[Ti(At(Date(2001, 3, 4); atol=Day(3)))] == 3
+            @test @inferred da[Ti(At(DateTime(2001, 3, 4); atol=Day(3)))] == 3
         end
         @testset "End locus" begin
             timedim = Ti(Sampled(DateTime(2001):Month(1):DateTime(2001, 12); 
@@ -1264,7 +1257,7 @@ end
         @test @inferred view(da, Between((11, 26)), Between((1.4u"s", 4u"s"))) == [6 7]
     end
 
-    @testset "with DateTime index" begin
+    @testset "with DateTime lookup" begin
         span_ti = Explicit(vcat(
             reshape((DateTime(2001, 1):Month(1):DateTime(2001, 12)), 1, 12),
             reshape((DateTime(2001, 2):Month(1):DateTime(2002, 1)), 1, 12)