From ff8824b6551cb8e15fec19bc4897e86944ebe394 Mon Sep 17 00:00:00 2001
From: Michael Abbott <me@escbook>
Date: Fri, 6 Mar 2020 13:30:52 -0700
Subject: [PATCH 1/3] make axes(A) return KeyedUnitRanges

---
 src/AxisKeys.jl |  2 ++
 src/axes.jl     | 52 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/struct.jl   |  2 +-
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 src/axes.jl

diff --git a/src/AxisKeys.jl b/src/AxisKeys.jl
index 0120b2d..c79e7db 100644
--- a/src/AxisKeys.jl
+++ b/src/AxisKeys.jl
@@ -3,6 +3,8 @@ module AxisKeys
 include("struct.jl")
 export KeyedArray, axiskeys
 
+include("axes.jl")
+
 include("lookup.jl")
 
 include("names.jl")
diff --git a/src/axes.jl b/src/axes.jl
new file mode 100644
index 0000000..6955f90
--- /dev/null
+++ b/src/axes.jl
@@ -0,0 +1,52 @@
+#=
+
+This is an experiment, with returning something very similar to a KeyedArray from axes(A).
+The reason to do so is that things like similar(C, axes(A,1), axes(B,2)) could propagate keys.
+However right now axes(A,1) !isa Base.OneTo results in lots of OffsetArrays...
+
+=#
+
+struct KeyedUnitRange{T,AT,KT} <: AbstractUnitRange{T}
+    data::AT
+    keys::KT
+    function KeyedUnitRange(data::AT, keys::KT) where {AT<:AbstractUnitRange{T}, KT<:AbstractVector} where {T}
+        new{T,AT,KT}(data, keys)
+    end
+end
+
+Base.parent(A::KeyedUnitRange) = getfield(A, :data)
+keyless(A::KeyedUnitRange) = parent(A)
+
+for f in [:size, :first, :last, :IndexStyle]
+    @eval Base.$f(A::KeyedUnitRange) = $f(parent(A))
+end
+
+Base.getindex(A::KeyedUnitRange, inds::Integer) = getindex(parent(A), inds)
+
+axiskeys(A::KeyedUnitRange) = tuple(getfield(A, :keys))
+axiskeys(A::KeyedUnitRange, d::Integer) = d==1 ? getfield(A, :keys) : Base.OneTo(1)
+
+
+# getkey(A::AbstractKeyedArray{<:Any,1}, key) = findfirst(isequal(key), axiskeys(A)[1])
+
+function Base.axes(A::KeyedArray)
+    ntuple(ndims(A)) do d
+        KeyedUnitRange(axes(parent(A),d), axiskeys(A, d))
+    end
+end
+
+KeyedUnion{T} = Union{KeyedArray{T}, KeyedUnitRange{T}}
+
+Base.summary(io::IO, x::KeyedUnion) = _summary(io, x)
+Base.summary(io::IO, A::NamedDimsArray{L,T,N,<:KeyedUnion}) where {L,T,N} = _summary(io, A)
+showtype(io::IO, ::KeyedUnitRange) = print(io, "KeyedUnitRange(...)")
+
+function Base.show(io::IO, m::MIME"text/plain", x::KeyedUnitRange)
+    summary(io, x)
+    println(io, ":")
+    keyed_print_matrix(io, x)
+end
+
+function Base.show(io::IO, x::KeyedUnitRange)
+    print(io, "KeyedUnitRange(", keyless(x), ", ", axiskeys(x,1), ")")
+end
diff --git a/src/struct.jl b/src/struct.jl
index b86591a..57de024 100644
--- a/src/struct.jl
+++ b/src/struct.jl
@@ -38,7 +38,7 @@ KeyedArray(A::KeyedVector, k2::Tuple{AbstractVector}) =
 
 Base.size(x::KeyedArray) = size(parent(x))
 
-Base.axes(x::KeyedArray) = axes(parent(x))
+# Base.axes(x::KeyedArray) = axes(parent(x))
 
 Base.parent(x::KeyedArray) = getfield(x, :data)
 keyless(x::KeyedArray) = parent(x)

From dcdbcf89037d6dce4a3f7b64471c89fa7183e4b9 Mon Sep 17 00:00:00 2001
From: Michael Abbott <me@escbook>
Date: Sat, 14 Mar 2020 15:52:10 -0400
Subject: [PATCH 2/3] load NamedDims

---
 src/axes.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/axes.jl b/src/axes.jl
index 6955f90..8456981 100644
--- a/src/axes.jl
+++ b/src/axes.jl
@@ -37,6 +37,8 @@ end
 
 KeyedUnion{T} = Union{KeyedArray{T}, KeyedUnitRange{T}}
 
+using NamedDims
+
 Base.summary(io::IO, x::KeyedUnion) = _summary(io, x)
 Base.summary(io::IO, A::NamedDimsArray{L,T,N,<:KeyedUnion}) where {L,T,N} = _summary(io, A)
 showtype(io::IO, ::KeyedUnitRange) = print(io, "KeyedUnitRange(...)")

From 3a55d437f4f17786856f1cebbd5e5a4cd0fdb7b2 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Wed, 6 Oct 2021 21:59:12 -0400
Subject: [PATCH 3/3] similar, reshape

---
 src/axes.jl    | 41 +++++++++++++++++++++++++++++++++++++----
 test/_basic.jl | 10 ++++++++++
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/src/axes.jl b/src/axes.jl
index 8456981..4a6455a 100644
--- a/src/axes.jl
+++ b/src/axes.jl
@@ -1,11 +1,18 @@
 #=
 
-This is an experiment, with returning something very similar to a KeyedArray from axes(A).
+This is an experiment, with returning something very similar to a KeyedArray from axes(A,1).
 The reason to do so is that things like similar(C, axes(A,1), axes(B,2)) could propagate keys.
 However right now axes(A,1) !isa Base.OneTo results in lots of OffsetArrays...
 
 =#
 
+"""
+    AxisKeys.KeyedUnitRange(Base.OneTo(3), ["a", "b", "c"])
+
+This exists to store both `axes(parent(A), d)`, and `axiskeys(A, d)` together,
+as one `AbstractUnitRange`, so that generic code which propagages `axes(A)`
+will pass keys along.
+"""
 struct KeyedUnitRange{T,AT,KT} <: AbstractUnitRange{T}
     data::AT
     keys::KT
@@ -17,24 +24,50 @@ end
 Base.parent(A::KeyedUnitRange) = getfield(A, :data)
 keyless(A::KeyedUnitRange) = parent(A)
 
-for f in [:size, :first, :last, :IndexStyle]
+for f in [:size, :length, :axes, :first, :last, :IndexStyle]
     @eval Base.$f(A::KeyedUnitRange) = $f(parent(A))
 end
 
-Base.getindex(A::KeyedUnitRange, inds::Integer) = getindex(parent(A), inds)
+Base.getindex(A::KeyedUnitRange, ind::Integer) = getindex(parent(A), ind)
 
 axiskeys(A::KeyedUnitRange) = tuple(getfield(A, :keys))
 axiskeys(A::KeyedUnitRange, d::Integer) = d==1 ? getfield(A, :keys) : Base.OneTo(1)
 
+haskeys(A::KeyedUnitRange) = true
 
 # getkey(A::AbstractKeyedArray{<:Any,1}, key) = findfirst(isequal(key), axiskeys(A)[1])
 
+
+# Use for KeyedArray, and for reconstruction of such.
+
 function Base.axes(A::KeyedArray)
     ntuple(ndims(A)) do d
-        KeyedUnitRange(axes(parent(A),d), axiskeys(A, d))
+        KeyedUnitRange(axes(parent(A), d), axiskeys(A, d))
     end
 end
 
+Base.similar(A::AbstractArray, ::Type{T}, ax::Tuple{KeyedUnitRange}) where {T} = _similar(A, T, ax)
+Base.similar(A::AbstractArray, ::Type{T}, ax::Tuple{AbstractUnitRange, KeyedUnitRange, Vararg{AbstractUnitRange}}) where {T} = _similar(A, T, ax)
+Base.similar(A::AbstractArray, ::Type{T}, ax::Tuple{KeyedUnitRange, AbstractUnitRange, Vararg{AbstractUnitRange}}) where {T} = _similar(A, T, ax)
+Base.similar(A::AbstractArray, ::Type{T}, ax::Tuple{KeyedUnitRange, KeyedUnitRange, Vararg{AbstractUnitRange}})  where {T} = _similar(A, T, ax)
+function _similar(A, ::Type{T}, ax) where {T}
+    data = similar(keyless(A), T, map(keyless, ax))
+    new_keys = map(a -> keys_or_axes(a,1), ax)
+    KeyedArray(data, new_keys)
+end
+
+Base.reshape(A::AbstractArray, ax::Tuple{KeyedUnitRange}) = _reshape(A, ax)
+Base.reshape(A::AbstractArray, ax::Tuple{AbstractUnitRange, KeyedUnitRange, Vararg{AbstractUnitRange}}) = _reshape(A, ax)
+Base.reshape(A::AbstractArray, ax::Tuple{KeyedUnitRange, AbstractUnitRange, Vararg{AbstractUnitRange}}) = _reshape(A, ax)
+Base.reshape(A::AbstractArray, ax::Tuple{KeyedUnitRange, KeyedUnitRange, Vararg{AbstractUnitRange}}) = _reshape(A, ax)
+function _reshape(A, ax)
+    data = reshape(keyless(A), map(keyless, ax))
+    new_keys = map(a -> keys_or_axes(a,1), ax)
+    KeyedArray(data, new_keys)
+end
+
+# Pretty printing
+
 KeyedUnion{T} = Union{KeyedArray{T}, KeyedUnitRange{T}}
 
 using NamedDims
diff --git a/test/_basic.jl b/test/_basic.jl
index 0a82245..013bbba 100644
--- a/test/_basic.jl
+++ b/test/_basic.jl
@@ -309,3 +309,13 @@ end
     @test axiskeys(W,1) == ["a", "b", "c", "d", "e"]
 
 end
+@testset "axes" begin
+    M = wrapdims(rand(Int8, 3,11), [:a, :b, :c], 0:0.1:1)
+    @test Base.axes(M) isa NTuple{2,AxisKeys.KeyedUnitRange}
+
+    # similar
+    @test axiskeys(similar(0:1, axes(M))) == axiskeys(M)
+    @test axiskeys(similar(0:1, ComplexF32, reverse(axes(M)))) == reverse(axiskeys(M))
+    @test_broken axiskeys(similar(rand(3), Base.OneTo(3), axes(M,2))) == (1:3, 0:0.1:1)
+
+end