make limited sparse printing O(1) (JuliaLang#22536)

* make limited sparse printing O(1)

* add test and print vdots when no element can be shown

* fix indent

base/sparse/sparsematrix.jl

@@ -135,32 +135,55 @@ function Base.show(io::IOContext, S::SparseMatrixCSC)
  if nnz(S) == 0
  return show(io, MIME("text/plain"), S)
  end
-
  limit::Bool = get(io, :limit, false)
- if limit
- rows = displaysize(io)[1]
- half_screen_rows = div(rows - 8, 2)
+ rows = displaysize(io)[1] - 4 # -4 from [Prompt, header, newline after elements, new prompt]
+ will_fit = !limit || rows >= nnz(S) # Will the whole matrix fit when printed?
+
+ if rows <= 2 && !will_fit
+ print(io, "\n \u22ee")
+ return
+ end
+
+ iob = IOBuffer()
+ ioc = IOContext(iob, :compact => true)
+
+ function _format_line(r, col)
+ pad = ndigits(max(S.m, S.n))
+ print(ioc, " ", '[', rpad(S.rowval[r], pad), ", ", lpad(col, pad), "] = ")
+ if isassigned(S.nzval, Int(r))
+ show(ioc, S.nzval[r])
+ else
+ print(ioc, Base.undef_ref_str)
+ end
+ return String(take!(iob))
+ end
+
+ if will_fit
+ print_count = nnz(S)
  else
- half_screen_rows = typemax(Int)
- end
- pad = ndigits(max(S.m,S.n))
- k = 0
- sep = "\n "
- if !haskey(io, :compact)
- io = IOContext(io, :compact => true)
- end
- for col = 1:S.n, k = S.colptr[col] : (S.colptr[col+1]-1)
- if k < half_screen_rows || k > nnz(S)-half_screen_rows
- print(io, sep, '[', rpad(S.rowval[k], pad), ", ", lpad(col, pad), "] = ")
- if isassigned(S.nzval, Int(k))
- show(io, S.nzval[k])
- else
- print(io, Base.undef_ref_str)
- end
- elseif k == half_screen_rows
- print(io, sep, '\u22ee')
+ print_count = div(rows-1, 2)
+ end
+
+ count = 0
+ for col = 1:S.n, r = nzrange(S, col)
+ count += 1
+ print(io, "\n", _format_line(r, col))
+ count == print_count && break
+ end
+
+ lines_bottom = String[]
+ if !will_fit
+ count = 0
+ for col = reverse(1:S.n), r = reverse(nzrange(S, col))
+ count += 1
+ push!(lines_bottom, _format_line(r, col))
+ count == print_count && break
  end
- k += 1
+ end
+
+ if !will_fit
+ print(io, "\n ", '\u22ee', "\n")
+ print(io, join(reverse(lines_bottom), '\n'))
  end
 end
 

test/sparse/sparse.jl

@@ -1772,6 +1772,26 @@ end
  @test String(take!(io)) == "1×1 SparseMatrixCSC{Float64,Int64} with 1 stored entry:\n [1, 1] = 1.0"
  show(io, MIME"text/plain"(), spzeros(Float32, Int64, 2, 2))
  @test String(take!(io)) == "2×2 SparseMatrixCSC{Float32,Int64} with 0 stored entries"
+
+ ioc = IOContext(IOContext(io, :displaysize => (5, 80)), :limit => true)
+ show(ioc, MIME"text/plain"(), sparse(Int64[1], Int64[1], [1.0]))
+ @test String(take!(io)) == "1×1 SparseMatrixCSC{Float64,Int64} with 1 stored entry:\n [1, 1] = 1.0"
+ show(ioc, MIME"text/plain"(), sparse(Int64[1, 1], Int64[1, 2], [1.0, 2.0]))
+ @test String(take!(io)) == "1×2 SparseMatrixCSC{Float64,Int64} with 2 stored entries:\n ⋮"
+
+ ioc = IOContext(IOContext(io, :displaysize => (9, 80)), :limit => true)
+ show(ioc, MIME"text/plain"(), sparse(Int64[1,2,3,4,5], Int64[1,1,2,2,3], [1.0,2.0,3.0,4.0,5.0]))
+ @test String(take!(io)) == string("5×3 SparseMatrixCSC{Float64,Int64} with 5 stored entries:\n [1, 1]",
+ " = 1.0\n [2, 1] = 2.0\n [3, 2] = 3.0\n [4, 2] = 4.0\n [5, 3] = 5.0")
+
+ show(ioc, MIME"text/plain"(), sparse(Int64[1,2,3,4,5,6], Int64[1,1,2,2,3,3], [1.0,2.0,3.0,4.0,5.0,6.0]))
+ @test String(take!(io)) == string("6×3 SparseMatrixCSC{Float64,Int64} with 6 stored entries:\n [1, 1]",
+ " = 1.0\n [2, 1] = 2.0\n ⋮\n [5, 3] = 5.0\n [6, 3] = 6.0")
+
+ ioc = IOContext(io, :displaysize => (9, 80))
+ show(ioc, MIME"text/plain"(), sparse(Int64[1,2,3,4,5,6], Int64[1,1,2,2,3,3], [1.0,2.0,3.0,4.0,5.0,6.0]))
+ @test String(take!(io)) == string("6×3 SparseMatrixCSC{Float64,Int64} with 6 stored entries:\n [1, 1] = 1.0\n",
+ " [2, 1] = 2.0\n [3, 2] = 3.0\n [4, 2] = 4.0\n [5, 3] = 5.0\n [6, 3] = 6.0")
 end
 
 @testset "similar aliasing" begin