Generalize printing (#102)

Author: mtfishman

Project.toml

@@ -1,6 +1,6 @@
 name = "KroneckerArrays"
 uuid = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
-version = "0.3.26"
+version = "0.3.27"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/kroneckerarray.jl

@@ -275,20 +275,26 @@ function Base.axes(ab::AbstractKroneckerArray)
     return axes(a) .× axes(b)
 end
 
-function Base.print_array(io::IO, ab::KroneckerArray)
+function Base.print_array(io::IO, ab::AbstractKroneckerArray)
     a, b = kroneckerfactors(ab)
     Base.print_array(io, a)
     println(io, "\n ⊗")
     Base.print_array(io, b)
     return nothing
 end
-function Base.show(io::IO, ab::KroneckerArray)
+function Base.show(io::IO, ab::AbstractKroneckerArray)
     a, b = kroneckerfactors(ab)
     show(io, a)
     print(io, " ⊗ ")
     show(io, b)
     return nothing
 end
+function Base.show(io::IO, ::MIME"text/plain", ab::AbstractKroneckerArray)
+    summary(io, ab)
+    println(io, ":")
+    Base.print_array(io, ab)
+    return nothing
+end
 
 ⊗(a1::AbstractArray, a2::AbstractArray) = KroneckerArray(a1, a2)
 ⊗(a1::Number, a2::Number) = a1 * a2

test/test_basics.jl

@@ -12,6 +12,20 @@ using StableRNGs: StableRNG
 using Test: @test, @test_broken, @test_throws, @testset
 using TestExtras: @constinferred
 
+module TestBasicsUtils
+    using KroneckerArrays: KroneckerArrays
+
+    struct WrappedKroneckerArray{T, N, A <: AbstractArray{T, N}, B <: AbstractArray{T, N}} <:
+        KroneckerArrays.AbstractKroneckerArray{T, N}
+        a::A
+        b::B
+    end
+    function WrappedKroneckerArray(a::AbstractArray{T, N}, b::AbstractArray{T, N}) where {T, N}
+        return WrappedKroneckerArray{T, N, typeof(a), typeof(b)}(a, b)
+    end
+    KroneckerArrays.kroneckerfactors(ab::WrappedKroneckerArray) = (ab.a, ab.b)
+end
+
 elts = (Float32, Float64, ComplexF32, ComplexF64)
 @testset "KroneckerArrays (eltype=$elt)" for elt in elts
     p = [1, 2] × [3, 4, 5]
@@ -141,6 +155,11 @@ elts = (Float32, Float64, ComplexF32, ComplexF64)
     @test bc.style === style
     @test collect(copy(bc)) ≈ 2 * collect(a)
 
+    # Display
+    aw = TestBasicsUtils.WrappedKroneckerArray(randn(elt, 2, 2), randn(elt, 3, 3))
+    @test occursin(" ⊗ ", sprint(show, aw))
+    @test occursin("\n ⊗\n", sprint(show, MIME("text/plain"), aw))
+
     # Mapping
     a = randn(elt, 2, 2) ⊗ randn(elt, 3, 3)
     @test_throws "not supported" map(sin, a)