Allow _AxisLookup to fallback to Dict getindex (#3028)

Author: nickrobinson251

src/Containers/DenseAxisArray.jl

@@ -48,17 +48,17 @@ function build_lookup(ax)
     return _AxisLookup(d)
 end
 
-Base.getindex(x::_AxisLookup{Dict{K,Int}}, key::K) where {K} = x.data[key]
+Base.getindex(x::_AxisLookup{Dict{K,Int}}, key) where {K} = x.data[key]
 Base.getindex(::_AxisLookup{Dict{K,Int}}, key::Colon) where {K} = key
 
 function Base.getindex(
     x::_AxisLookup{Dict{K,Int}},
-    keys::AbstractVector{<:K},
+    keys::AbstractVector,
 ) where {K}
     return [x[key] for key in keys]
 end
 
-function Base.get(x::_AxisLookup{Dict{K,Int}}, key::K, default) where {K}
+function Base.get(x::_AxisLookup{Dict{K,Int}}, key, default) where {K}
     return get(x.data, key, default)
 end
 

test/Containers/DenseAxisArray.jl

@@ -43,6 +43,7 @@ using Test
         @test isassigned(A, 2)
         @test !isassigned(A, 1)
         @test length.(axes(A)) == (2,)
+        @test_throws KeyError A["2"]
 
         correct_answer = DenseAxisArray([2.0, 3.0], 2:3)
         @test sprint(show, correct_answer) == """
@@ -93,6 +94,14 @@ And data, a 2-element $(Vector{Float64}):
         @test 1 .+ A == correct_answer
     end
 
+    @testset "String index set" begin
+        A = @inferred DenseAxisArray([1.0, 2.0], ["a", "b"])
+        @test (@inferred A["a"]) == (@inferred A[GenericString("a")]) == 1.0
+        @test (@inferred A[["a", "b"]]) ==
+              (@inferred A[[GenericString("a"), GenericString("b")]]) ==
+              A
+    end
+
     @testset "Mixed range/symbol index sets" begin
         A = @inferred DenseAxisArray([1 2; 3 4], 2:3, [:a, :b])
         @test size(A) == (2, 2)
@@ -248,6 +257,11 @@ And data, a 0-dimensional $(Array{Int,0}):
         @test (@inferred C[2:3, [:a, :b]]) == C
         @test (@inferred C[2, [:a, :b]]) == DenseAxisArray([5.0, 6.0], [:a, :b])
         @test (@inferred C[2:3, :b]) == DenseAxisArray([6.0, 8.0], 2:3)
+
+        D = DenseAxisArray([5.0 6.0; 7.0 8.0], 2:3, ["a", "b"])
+        @test (@inferred D[2, GenericString("b")]) == 6.0
+        @test (@inferred D[2, [GenericString("a"), GenericString("b")]]) ==
+              DenseAxisArray([5.0, 6.0], ["a", "b"])
     end
     @testset "BitArray" begin
         x = DenseAxisArray([0 1; 1 0], [:a, :b], 1:2)