Add prison_count

src/Tally.jl

@@ -4,7 +4,7 @@ import Printf: @sprintf
 
 import UnicodePlots: barplot, label!
 
-export tally, lazy_tally, materialize
+export tally, lazy_tally, materialize, prison_count, show_style
 
 ################################################################################
 #
@@ -209,8 +209,7 @@ function _get_nice_percentages(per::Vector{Float64})
   return find_ndigits, nice
 end
 
-# Print it as a table
-function Base.show(io::IO, ::MIME"text/plain", T::TallyT)
+function _show_as_table(io::IO, T::TallyT)
   first = true
   n = sum(T.values)
   k, v = _prepare_for_plot(T)
@@ -232,6 +231,87 @@ function Base.show(io::IO, ::MIME"text/plain", T::TallyT)
   end
 end
 
+function _show_as_plot(io::IO, T::TallyT)
+  print(io, plot(T))
+end
+
+function _prison_count(n)
+  q, r = divrem(n, 5)
+  return "┼┼┼┼ "^q * "│"^r
+end
+
+function _show_as_prison(io::IO, T::TallyT)
+  k, v = _prepare_for_plot(T)
+  l_names, l_digits = _maximal_length_of_items(k, v)
+  m = isempty(values(T)) ? 0 : maximum(values(T))
+  first = true
+  for (kk, vv) in zip(keys(T), values(T))
+    if !first
+      println(io, "━"^(l_names + 1) * "╋" * "━"^(m + 1))
+    else
+      first = false
+    end
+    print(io, rpad(kk, l_names), " ┃" )
+    print(io, " ", _prison_count(vv))
+    println(io)
+  end
+end
+
+"""
+    prison_count(T::Tally)
+
+Print the tally using prison style:
+```jldoctest
+julia> T = tally([1, 1, 1, 2, 2, 2, 2, 2, 2, -1]);
+
+julia> prison_count(T)
+2  ┃ ┼┼┼┼ │
+━━━╋━━━━━━━
+1  ┃ │││
+━━━╋━━━━━━━
+-1 ┃ │
+```
+"""
+function prison_count(T::TallyT)
+  _show_as_prison(stdout, T)
+end
+
+SHOW_STYLE = Ref(:table)
+
+function show_style()
+  SHOW_STYLE[]
+end
+
+"""
+    show_style([style::Symbol])
+
+With no argument, this function returns the current style used for the `show` method,
+that is, for printing.
+
+If `style::Symbol` is provided, this sets the show style to `style`. Possible styles are
+- `:table`, the default,
+- `:plot`, print tallies via `Tally.plot`,
+- `:prison`, display the counts using prison mode.
+"""
+function show_style(style::Symbol)
+  if !(style in [:table, :plot, :prison])
+    error("Style must be either, `:table`, `:plot` or `:prison`")
+  end
+  SHOW_STYLE[] = style
+end
+
+# Print it as a table
+function Base.show(io::IO, ::MIME"text/plain", T::TallyT)
+  S = show_style()
+  if S === :table
+    _show_as_table(io, T)
+  elseif S === :plot
+    _show_as_plot(io, T)
+  elseif S === :prison
+    _show_as_prison(io, T)
+  end
+end
+
 # Print it inline
 function Base.show(io::IO, T::TallyT)
   print(io, "Tally(")

test/runtests.jl

@@ -58,31 +58,55 @@ Base.hash(::A) = rand(UInt)
   @test T.keys == [A(1)]
   @test T.values == [3]
 
+  @test_throws ErrorException show_style(:bla)
 
   T = tally([2, 1, 1, 1, 1, 2, 2, 3])
-
   @test sprint(show, "text/plain", T) isa String
   @test sprint(show, T) isa String
+  prison_count(T)
 
   T = tally([2, 1, 1, 1, 1, 2, 2, 3], by = x -> x^2)
   s = sprint(show, "text/plain", T)
   @test occursin("[1]", s)
+  prison_count(T)
+  show_style(:prison)
+  s = sprint(show, "text/plain", T)
+  @test occursin("│", s)
+  show_style(:table)
+
 
   T = tally([2, 1, 1, 1, 1, 2, 2, 3], equivalence = (x, y) -> x == y)
   s = sprint(show, "text/plain", T)
   @test occursin("[1]", s)
+  prison_count(T)
+  show_style(:prison)
+  s = sprint(show, "text/plain", T)
+  @test occursin("│", s)
+  show_style(:table)
 
   T = tally(Int[])
   @test sprint(show, "text/plain", T) isa String
   @test sprint(show, T) isa String
+  prison_count(T)
+  show_style(:prison)
+  s = sprint(show, "text/plain", T)
+  show_style(:table)
 
   T = tally(push!([1 for i in 1:100000], 2))
   s = sprint(show, "text/plain", T)
   @test occursin("1", s)
+  show_style(:prison)
+  s = sprint(show, "text/plain", T)
+  @test occursin("│", s)
+  show_style(:table)
 
   T = tally(append!([1 for i in 1:100000], [2 for i in 1:100001]))
   s = sprint(show, "text/plain", T)
   @test occursin("1", s)
+  show_style(:prison)
+  s = sprint(show, "text/plain", T)
+  @test occursin("│", s)
+  show_style(:table)
 
   # unicode plotting
 
@@ -146,6 +170,11 @@ Base.hash(::A) = rand(UInt)
     "     └                                        ┘       "
 
     @test string(Tally.plot(T)) == plot1
+    show_style(:plot)
+    s = sprint(show, "text/plain", T)
+    @test s == plot1
+    show_style(:table)
+
     @test string(Tally.plot(T, sortby = :value)) == plot1
     @test string(Tally.plot(T, sortby = :key)) == plot2
     @test string(Tally.plot(T, sortby = :key, reverse = true)) == plot1