Merge pull request #772 from dbcli/koljonen/array_formatting

Better formatting for array output

4 files changed, 103 insertions, 6 deletions

diff --git a/changelog.rst b/changelog.rst
index 36cba819..98050104 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -5,6 +5,7 @@ Features:
 ---------
 
 * Add fish-style auto-suggestion from history. (Thanks: `Amjith Ramanujam`_)
+* Improved formatting of arrays in output (Thanks: `Joakim Koljonen`_)
 * Don't quote identifiers that are non-reserved keywords. (Thanks: `Joakim Koljonen`_)
 * Remove the ``...`` in the continuation prompt and use empty space instead. (Thanks: `Amjith Ramanujam`_)
 
diff --git a/pgcli/encodingutils.py b/pgcli/encodingutils.py
index adc4502b..5de97b7e 100644
--- a/pgcli/encodingutils.py
+++ b/pgcli/encodingutils.py
@@ -3,6 +3,8 @@ import sys
 PY2 = sys.version_info[0] == 2
 PY3 = sys.version_info[0] == 3
 
+text_type = unicode if PY2 else str
+
 
 def unicode2utf8(arg):
     """
diff --git a/pgcli/main.py b/pgcli/main.py
index 45149f8f..5bd883b5 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -50,6 +50,7 @@ from .config import (get_casing_file,
     load_config, config_location, ensure_dir_exists, get_config)
 from .key_bindings import pgcli_bindings
 from .encodingutils import utf8tounicode
+from .encodingutils import text_type
 from .__init__ import __version__
 
 click.disable_unicode_literals_warning = True
@@ -64,6 +65,7 @@ from psycopg2 import OperationalError, InterfaceError
 
 from collections import namedtuple
 
+
 # Query tuples are used for maintaining history
 MetaQuery = namedtuple(
     'Query',
@@ -946,6 +948,23 @@ def format_output(title, cur, headers, status, settings):
     case_function = settings.case_function
     formatter = TabularOutputFormatter(format_name=table_format)
 
+    def format_array(val):
+        if val is None:
+            return settings.missingval
+        if not isinstance(val, list):
+            return val
+        return '{' + ','.join(text_type(format_array(e)) for e in val) + '}'
+
+    def format_arrays(data, headers, **_):
+        data = list(data)
+        for row in data:
+            row[:] = [
+                format_array(val) if isinstance(val, list) else val
+                for val in row
+            ]
+
+        return data, headers
+
     output_kwargs = {
         'sep_title': 'RECORD {n}',
         'sep_character': '-',
@@ -953,7 +972,7 @@ def format_output(title, cur, headers, status, settings):
         'missing_value': settings.missingval,
         'integer_format': settings.dcmlfmt,
         'float_format': settings.floatfmt,
-        'preprocessors': (format_numbers, ),
+        'preprocessors': (format_numbers, format_arrays),
         'disable_numparse': True,
         'preserve_whitespace': True
     }
diff --git a/tests/test_main.py b/tests/test_main.py
index 0380c6fa..c35853ee 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,6 +1,9 @@
+# coding=utf-8
+from __future__ import unicode_literals
 import os
 import platform
 import mock
+from decimal import Decimal
 
 import pytest
 try:
@@ -52,7 +55,61 @@ def test_format_output():
     settings = OutputSettings(table_format='psql', dcmlfmt='d', floatfmt='g')
     results = format_output('Title', [('abc', 'def')], ['head1', 'head2'],
                             'test status', settings)
-    expected = ['Title', '+---------+---------+\n| head1   | head2   |\n|---------+---------|\n| abc     | def     |\n+---------+---------+', 'test status']
+    expected = [
+        'Title',
+        '+---------+---------+\n'
+        '| head1   | head2   |\n'
+        '|---------+---------|\n'
+        '| abc     | def     |\n'
+        '+---------+---------+',
+        'test status'
+    ]
+    assert results == expected
+
+
+def test_format_array_output(executor):
+    statement = u"""
+    SELECT
+        array[1, 2, 3]::bigint[] as bigint_array,
+        '{{1,2},{3,4}}'::numeric[] as nested_numeric_array,
+        '{å,魚,текст}'::text[] as 配列
+    UNION ALL
+    SELECT '{}', NULL, array[NULL]
+    """
+    results = run(executor, statement)
+    expected = [
+        '+----------------+------------------------+--------------+\n'
+        '| bigint_array   | nested_numeric_array   | 配列         |\n'
+        '|----------------+------------------------+--------------|\n'
+        '| {1,2,3}        | {{1,2},{3,4}}          | {å,魚,текст} |\n'
+        '| {}             | <null>                 | {<null>}     |\n'
+        '+----------------+------------------------+--------------+',
+        'SELECT 2'
+    ]
+    assert results == expected
+
+
+def test_format_array_output_expanded(executor):
+    statement = u"""
+    SELECT
+        array[1, 2, 3]::bigint[] as bigint_array,
+        '{{1,2},{3,4}}'::numeric[] as nested_numeric_array,
+        '{å,魚,текст}'::text[] as 配列
+    UNION ALL
+    SELECT '{}', NULL, array[NULL]
+    """
+    results = run(executor, statement, expanded=True)
+    expected = [
+        '-[ RECORD 1 ]-------------------------\n'
+        'bigint_array         | {1,2,3}\n'
+        'nested_numeric_array | {{1,2},{3,4}}\n'
+        '配列                   | {å,魚,текст}\n'
+        '-[ RECORD 2 ]-------------------------\n'
+        'bigint_array         | {}\n'
+        'nested_numeric_array | <null>\n'
+        '配列                   | {<null>}\n',
+        'SELECT 2'
+    ]
     assert results == expected
 
 
@@ -61,12 +118,30 @@ def test_format_output_auto_expand():
         table_format='psql', dcmlfmt='d', floatfmt='g', max_width=100)
     table_results = format_output('Title', [('abc', 'def')],
                                   ['head1', 'head2'], 'test status', settings)
-    table = ['Title', '+---------+---------+\n| head1   | head2   |\n|---------+---------|\n| abc     | def     |\n+---------+---------+', 'test status']
+    table = [
+        'Title',
+        '+---------+---------+\n'
+        '| head1   | head2   |\n'
+        '|---------+---------|\n'
+        '| abc     | def     |\n'
+        '+---------+---------+',
+        'test status'
+    ]
     assert table_results == table
-    expanded_results = format_output('Title', [('abc', 'def')],
-                                     ['head1', 'head2'], 'test status', settings._replace(max_width=1))
+    expanded_results = format_output(
+        'Title',
+        [('abc', 'def')],
+        ['head1', 'head2'],
+        'test status',
+        settings._replace(max_width=1)
+    )
     expanded = [
-        'Title', u'-[ RECORD 1 ]-------------------------\nhead1 | abc\nhead2 | def\n', 'test status']
+        'Title',
+        '-[ RECORD 1 ]-------------------------\n'
+        'head1 | abc\n'
+        'head2 | def\n',
+        'test status'
+    ]
     assert expanded_results == expanded