improve printing of GroupedDataFrame in corner cases (#3187)

Author: bkamins

src/groupeddataframe/show.jl

@@ -42,30 +42,42 @@ function Base.show(io::IO, gd::GroupedDataFrame;
 
         (h, w) = displaysize(io)
 
+        # in the code below we accept that output for desired height less
+        # than 15 does not have to always exactly match the passed value
         if h > 0
-            h -= 2 # two lines are already used for header and gap between groups
-            h <= 5 && (h = 5) # if too small to fit, print fully compact
-        end
-
-        h1 = h2 = h # display heights available for first and last groups
-        if N > 1 && h > 0
-
-            # line height of groups if printed in full (nrows + 3 extra for header)
-            g1 = size(gd[1], 1) + 3
-            g2 = size(gd[N], 1) + 3
-
-            if g1 + g2 > h # won't fit on screen
-                if g1 < h ÷ 2
-                    h2 = h - g1 - 2 # show first group fully, squash last
-                elseif g2 < h ÷ 2
-                    h1 = h - g2 - 2 # show last group fully, squash first
+            # 2 lines for header and gap between groups, 3 lines for prompts;
+            # correcting for this allows at least 8 lines (4 for each group)
+            h = max(h - 5, 8)
+
+            if N == 1
+                h1 = h + 3 # add two lines for prompts and one line as there is no gap between groups
+                h2 = 0 # not used
+            else
+                # line height of groups if printed in full; 4 lines for header
+                # we assume scenario where eltype is printed for simplicity
+                g1 = size(gd[1], 1) + 4
+                g2 = size(gd[N], 1) + 4
+                # below +2 is for 2 lines for prompts as we do not print summary
+                if g1 + g2 > h # won't fit on screen
+                    if g1 <= h ÷ 2
+                        h1 = g1 + 2
+                        h2 = h - g1 + 2 # show first group fully, squash last
+                    elseif g2 <= h ÷ 2
+                        h1 = h - g2 + 2 # show last group fully, squash first
+                        h2 = g2 + 2
+                    else
+                        # squash both groups
+                        h2 = h ÷ 2 + 2
+                        h1 = h - h2 + 4
+                    end
                 else
-                    # squash both groups
-                    h += 1
-                    h2 = h ÷ 2
-                    h1 = h - h2
+                    h1 = g1 + 2
+                    h2 = g2 + 2
                 end
             end
+        else
+            h1 = h # no limit
+            h2 = h # no limit
         end
 
         nrows = size(gd[1], 1)

test/show.jl

@@ -261,7 +261,7 @@ end
            4 │ true  true  true
            5 │ true  true  true"""
 
-    show(io, groupby(df, :x), allcols=true, allrows=true) 
+    show(io, groupby(df, :x), allcols=true, allrows=true)
     str = String(take!(io.io))
     @test str == """
         GroupedDataFrame with 2 groups based on key: x
@@ -336,22 +336,35 @@ end
     end
 
     # printed height always matches desired height, above a reasonable minimum
-    for h in 15:40
+    for a in 1:50, b in 1:50, h in 15:40
+        df = DataFrame(x = [fill(1, a); fill(2, b)])
         io = IOContext(IOBuffer(), :displaysize=>(h, 40), :limit=>true)
         show(io, groupby(df, :x), allcols=true)
         str = String(take!(io.io))
         nlines = length(split(str, '\n'))
-        desired = h - 3 # leave one line for last REPL prompt at top, two for new prompt
+        # leave one line for last REPL prompt at top, two for new prompt
         # (this is the same behavior as ungrouped data frames)
+        desired = min(a + b + 10, h - 3)
+        @test nlines == desired
+    end
+
+    for a in 1:50, h in 15:40
+        df = DataFrame(x = fill(1, a))
+        io = IOContext(IOBuffer(), :displaysize=>(h, 40), :limit=>true)
+        show(io, groupby(df, :x), allcols=true)
+        str = String(take!(io.io))
+        nlines = length(split(str, '\n'))
+        # leave one line for last REPL prompt at top, two for new prompt
+        # (this is the same behavior as ungrouped data frames)
+        desired = min(a + 5, h - 3)
         @test nlines == desired
     end
 
     # one group
     io = IOContext(IOBuffer(), :displaysize=>(15, 40), :limit=>true)
-    df = DataFrame(x = 1:15, y = 1)
+    df = DataFrame(x = Int64.(1:15), y = Int64(1))
     show(io, groupby(df, :y))
     str = String(take!(io.io))
-    print(str)
     @test str == """
         GroupedDataFrame with 1 group based on key: y
         First Group (15 rows): y = 1