From 1b8a8100befddb5108790bffd95738d3353bfcde Mon Sep 17 00:00:00 2001
From: Jonathan Doucette <jdoucette@physics.ubc.ca>
Date: Wed, 27 Nov 2024 15:10:28 -0800
Subject: [PATCH] extend `KeyPath` to allow `CartesianIndex` keys

---
 src/keypath.jl  |  9 ++++++---
 test/keypath.jl | 13 +++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/keypath.jl b/src/keypath.jl
index d6f0b17..2cadc1d 100644
--- a/src/keypath.jl
+++ b/src/keypath.jl
@@ -1,18 +1,18 @@
 using Base: tail
 
-KeyT = Union{Symbol, AbstractString, Integer}
+KeyT = Union{Symbol, AbstractString, Integer, CartesianIndex}
 
 """
     KeyPath(keys...)
 
 A type for representing a path of keys to a value in a nested structure.
 Can be constructed with a sequence of keys, or by concatenating other `KeyPath`s.
-Keys can be of type `Symbol`, `String`, or `Int`.
+Keys can be of type `Symbol`, `String`, `Int`, or `CartesianIndex`.
 
 For custom types, access through symbol keys is assumed to be done with `getproperty`.
 For consistency, the method `Base.propertynames` is used to get the viable property names.
 
-For string and integer keys instead, the access is done with `getindex`.
+For string, integer, and cartesian index keys, the access is done with `getindex` instead.
 
 See also [`getkeypath`](@ref), [`haskeypath`](@ref).
 
@@ -85,11 +85,13 @@ end
 keypathstr(kp::KeyPath) = join(kp.keys, ".")
 
 _getkey(x, k::Integer) = x[k]
+_getkey(x::AbstractArray, k::CartesianIndex) = x[k]
 _getkey(x, k::Symbol) = getproperty(x, k)
 _getkey(x::AbstractDict, k::Symbol) = x[k]
 _getkey(x, k::AbstractString) = x[k]
 
 _setkey!(x, k::Integer, v) = (x[k] = v)
+_setkey!(x::AbstractArray, k::CartesianIndex, v) = (x[k] = v)
 _setkey!(x, k::Symbol, v) = setproperty!(x, k, v)
 _setkey!(x::AbstractDict, k::Symbol, v) = (x[k] = v)
 _setkey!(x, k::AbstractString, v) = (x[k] = v)
@@ -97,6 +99,7 @@ _setkey!(x, k::AbstractString, v) = (x[k] = v)
 _haskey(x, k::Integer) = haskey(x, k)
 _haskey(x::Tuple, k::Integer) = 1 <= k <= length(x)
 _haskey(x::AbstractArray, k::Integer) = 1 <= k <= length(x) # TODO: extend to generic indexing
+_haskey(x::AbstractArray, k::CartesianIndex) = checkbounds(Bool, x, k)
 _haskey(x, k::Symbol) = k in propertynames(x)
 _haskey(x::AbstractDict, k::Symbol) = haskey(x, k)
 _haskey(x, k::AbstractString) = haskey(x, k)
diff --git a/test/keypath.jl b/test/keypath.jl
index a303a02..6fa648c 100644
--- a/test/keypath.jl
+++ b/test/keypath.jl
@@ -44,6 +44,10 @@
         kp = KeyPath(:b, :c, 2)
         @test getkeypath(x, kp) == 7
 
+        x = [(a=1,) (b=2,)]
+        @test getkeypath(x, KeyPath(CartesianIndex(1, 1), :a)) == 1
+        @test getkeypath(x, KeyPath(CartesianIndex(1, 2), :b)) == 2
+
         @testset "access through getproperty" begin
             x = Tkp(3, Dict(:c => 4, :d => 5), 6);
 
@@ -65,6 +69,10 @@
         kp = KeyPath(:b, :c, 2)
         setkeypath!(x, kp, 17)
         @test x.b.c[2] == 17
+
+        x = [(a=1,) (b=2,)]
+        setkeypath!(x, KeyPath(CartesianIndex(1, 2)), (c=3,))
+        @test x[2] == (c=3,)
     end
 
     @testset "haskeypath" begin
@@ -75,6 +83,11 @@
         @test !haskeypath(x, KeyPath(:b, "d", 4))
         @test !haskeypath(x, KeyPath(:b, "e"))
 
+        x = [(a=1,) (b=2,)]
+        @test haskeypath(x, KeyPath(CartesianIndex(1, 1)))
+        @test haskeypath(x, KeyPath(CartesianIndex(1, 2)))
+        @test !haskeypath(x, KeyPath(CartesianIndex(1, 3)))
+
         @testset "access through getproperty" begin
             x = Tkp(3, Dict(:c => 4, :d => 5), 6);