print the grid in terminal, modified based on PR #2

README.md

@@ -114,6 +114,21 @@ When drawing tasks, arckit will intelligently resize all of the grids such that
 
 ![Example of arckit output](./images/arcsave_example.png)
 
+Alternatively, the `print_grid` function outputs a grid directly to the terminal without creating any files. Note that it takes a numpy array as input, which is the standard grid format for a *Task*.
+
+```python
+>>> task = train_set[0] # Get the first task
+>>> for i, (input_array, output_array) in enumerate(task.train): # Loop through the training examples
+>>>     print(f"Training Example {i+1}")
+>>>     print("Input:")
+>>>     vis.print_grid(input_array)
+>>>     print("Output:")
+>>>     vis.print_grid(output_array)
+>>>     print()
+```
+
+![Example of print grid in terminal](./images/arc_terminal_print_grid_example.png)
+
 ## 💻 Command-line tools
 
 `arcshow` draws a visualisation of a specific task straight to the console:

arckit/vis.py

@@ -1,6 +1,7 @@
 import drawsvg
 import numpy as np
 import io
+import rich
 
 cmap = [
     '#252525', # black
@@ -262,4 +263,45 @@ def output_drawing(d: drawsvg.Drawing, filename: str, context=None):
         import cairosvg
         cairosvg.svg2pdf(bytestring=buffer.getvalue(), write_to=filename)
     else:
-        raise ValueError(f'Unknown file extension for {filename}')
+        raise ValueError(f'Unknown file extension for {filename}')
+
+def print_grid(grid: np.ndarray):
+    """
+    Print a grid to the terminal using rich library
+
+    Parameters
+    ----------
+    grid : np.ndarray
+        the standard grid format for Task 
+    """
+    CELL_WIDTH = 2
+
+    def get_color(color_str):
+        color_str = color_str.strip('#')
+        return rich.color.Color.from_rgb(*bytes.fromhex(color_str))
+
+    # Translate 'cmap' to rich style with respective color as background
+    rich_scmap = {
+        i: rich.style.Style(bgcolor=get_color(color))
+        for i, color in enumerate(cmap)
+    }
+
+    # Create and populate rich table
+    table = rich.table.Table.grid(expand=False)
+
+    height, width = grid.shape
+
+    # Add columns
+    for x in range(width):
+        table.add_column()
+        table.columns[x]._cells = [''] * height
+
+    table.rows = [rich.table.Row()] * height
+
+    # Populate cells
+    for y in range(height):
+        for x in range(width):
+            color_idx = grid[y, x]
+            table.columns[x]._cells[y] = rich.text.Text(' ' * CELL_WIDTH, style=rich_scmap[color_idx])
+
+    rich.print(table)