diag for Eye (#270)

* diag for Eye * namespace
JuliaArrays · Jul 3, 2023 · ac19566 · ac19566
1 parent 1ba83ba
commit ac19566
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/src/FillArrays.jl b/src/FillArrays.jl
@@ -434,6 +434,12 @@ const Eye{T,Axes} = RectOrDiagonal{T,Ones{T,1,Tuple{Axes}}}
 
 isone(::SquareEye) = true
 
+function diag(E::Eye, k::Integer=0)
+    v = k == 0 ? oneunit(eltype(E)) : zero(eltype(E))
+    len = length(diagind(E, k))
+    Fill(v, len)
+end
+
 # These should actually be in StdLib, LinearAlgebra.jl, for all Diagonal
 for f in (:permutedims, :triu, :triu!, :tril, :tril!, :copy)
     @eval ($f)(IM::Diagonal{<:Any,<:AbstractFill}) = IM

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -470,6 +470,13 @@ end
     @test Diagonal(Eye(8,5)) == Diagonal(ones(5))
     @test convert(Diagonal, Eye(5)) == Diagonal(ones(5))
     @test convert(Diagonal{Int}, Eye(5)) == Diagonal(ones(Int,5))
+
+    for E in (Eye(2,4), Eye(3))
+        M = collect(E)
+        for i in -5:5
+            @test diag(E, i) == diag(M, i)
+        end
+    end
 end
 
 @testset "one" begin