[hail] Teach Expressions to HTML Show

hail/python/hail/expr/expressions/base_expression.py

@@ -593,6 +593,9 @@ def __ne__(self, other):
         return self._compare_op("!=", other)
 
     def _to_table(self, name):
+        return self._to_relational(self, name, force_table=True)
+
+    def _to_relational(self, name, force_table=False):
         source = self._indices.source
         axes = self._indices.axes
         if not self._aggregations.empty():
@@ -650,15 +653,17 @@ def _to_table(self, name):
             assert len(axes) == 2
             assert isinstance(source, hail.MatrixTable)
             source = source.select_entries(**{name: self}).select_rows().select_cols()
-            to_return = source.key_cols_by().entries().select_globals()
+            if force_table:
+                to_return = source.key_cols_by().entries().select_globals()
+            else:
+                to_return = source.select_globals()
         assert self.dtype == to_return[name].dtype, f'type mismatch:\n' \
                                                     f'  Actual:    {self.dtype}\n' \
                                                     f'  Should be: {to_return[name].dtype}'
         return to_return
 
-
-    @typecheck_method(n=int, width=int, truncate=nullable(int), types=bool, handler=anyfunc)
-    def show(self, n=10, width=90, truncate=None, types=True, handler=print):
+    @typecheck_method(n=int, width=int, truncate=nullable(int), types=bool, handler=nullable(anyfunc))
+    def show(self, n=10, width=90, truncate=None, types=True, handler=None):
         """Print the first few rows of the table to the console.
 
         Examples
@@ -701,33 +706,34 @@ def show(self, n=10, width=90, truncate=None, types=True, handler=print):
         types : :obj:`bool`
             Print an extra header line with the type of each field.
         """
-        handler(self._show(n, width, truncate, types))
+        self._to_relational_preserving_rows_and_cols().show(
+            n_rows=n, width=width, truncate=truncate, types=types, handler=handler)
 
-    def _show(self, n, width, truncate, types):
+    def _to_relational_preserving_rows_and_cols(self):
         name = '<expr>'
         source = self._indices.source
         if isinstance(source, hl.Table):
             if self is source.row:
-                return source._show(n, width, truncate, types)
+                return source
             elif self is source.key:
-                return source.select()._show(n, width, truncate, types)
+                return source.select()
         elif isinstance(source, hl.MatrixTable):
             if self is source.row:
-                return source.rows()._show(n, width, truncate, types)
+                return source.rows()
             elif self is source.row_key:
-                return source.rows().select()._show(n, width, truncate, types)
+                return source.rows().select()
             if self is source.col:
-                return source.cols()._show(n, width, truncate, types)
+                return source.cols()
             elif self is source.col_key:
-                return source.cols().select()._show(n, width, truncate, types)
+                return source.cols().select()
             if self is source.entry:
-                return source.select_rows().select_cols().entries()._show(n, width, truncate, types)
+                return source.select_rows().select_cols()
         if source is not None:
             name = source._fields_inverse.get(self, name)
-        t = self._to_table(name)
-        if name in t.key:
-            t = t.order_by(*t.key).select(name)
-        return t._show(n, width, truncate, types)
+        x = self._to_relational(name)
+        if isinstance(x, hl.Table) and name in x.key:
+            return x.order_by(*x.key).select(name)
+        return x
 
 
     @typecheck_method(n=int, _localize=bool)

hail/python/hail/table.py

@@ -1396,8 +1396,8 @@ def format_line(values):
 
         return s
 
-    @typecheck_method(n=nullable(int), width=nullable(int), truncate=nullable(int), types=bool, handler=nullable(anyfunc))
-    def show(self, n=None, width=None, truncate=None, types=True, handler=None):
+    @typecheck_method(n=nullable(int), width=nullable(int), truncate=nullable(int), types=bool, handler=nullable(anyfunc), n_rows=nullable(int))
+    def show(self, n=None, width=None, truncate=None, types=True, handler=None, n_rows=None):
         """Print the first few rows of the table to the console.
 
         Examples
@@ -1418,7 +1418,7 @@ def show(self, n=None, width=None, truncate=None, types=True, handler=None):
 
         Parameters
         ----------
-        n : :obj:`int`
+        n or n_rows : :obj:`int`
             Maximum number of rows to show.
         width : :obj:`int`
             Horizontal width at which to break fields.
@@ -1430,6 +1430,11 @@ def show(self, n=None, width=None, truncate=None, types=True, handler=None):
         handler : Callable[[str], Any]
             Handler function for data string.
         """
+        if n_rows is not None and n is not None:
+            raise ValueError(f'specify one of n_rows or n, recieved {n_rows} and {n}')
+        if n_rows is not None:
+            n = n_rows
+        del n_rows
         if handler is None:
             try:
                 from IPython.display import display

hail/python/test/hail/expr/test_show.py

@@ -0,0 +1,24 @@
+from ..helpers import startTestHailContext, stopTestHailContext
+import unittest.TestCase
+
+import hail as hl
+
+setUpModule = startTestHailContext
+tearDownModule = stopTestHailContext
+
+class Tests(unittest.TestCase):
+    def test():
+        mt = hl.balding_nichols_model(3, 10, 10)
+        t = mt.rows()
+        mt.GT.show()
+        mt.locus.show()
+        mt.af.show()
+        mt.pop.show()
+        mt.sample_idx.show()
+        mt.bn.show()
+        mt.bn.fst.show()
+        mt.GT.n_alt_alleles().show()
+        (mt.GT.n_alt_alleles() * mt.GT.n_alt_alleles()).show()
+        (mt.af * mt.GT.n_alt_alleles()).show()
+        t.af.show()
+        (t.af * 3).show()