Respect \pset pager on expected behavior

"\pset pager" has three possible values: "always", "on", and "off".
pgcli previously treated all non-"off" values as "always". This change
implements the expected behavior, which is to use the pager when the
output is larger than the terminal height (See \pset pager in
https://www.postgresql.org/docs/9.2/static/app-psql.html). Pgcli adds to
this by also using the pager when the output is wider than the terminal
width.

Fixes #813

7 files changed, 125 insertions, 4 deletions

diff --git a/.gitignore b/.gitignore
index 55b4a74f..170585df 100644
--- a/.gitignore
+++ b/.gitignore
@@ -66,3 +66,6 @@ target/
 # Generated Packages
 *.deb
 *.rpm
+
+.vscode/
+venv/
diff --git a/AUTHORS b/AUTHORS
index dff41f7c..96b4c221 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -84,6 +84,7 @@ Contributors:
     * Saif Hakim
     * Artur Balabanov
     * Kenny Do
+    * Max Rothman
 
 
 Creator:
diff --git a/changelog.rst b/changelog.rst
index 90665930..4cebf27c 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -1,3 +1,11 @@
+Upcoming
+========
+
+Features:
+---------
+
+* Respect `\pset pager on` and use pager when output is longer than terminal height (Thanks: `Max Rothman`_)
+
 1.10.3
 ======
 
@@ -868,3 +876,4 @@ Improvements:
 .. _`Saif Hakim`: https://github.com/saifelse
 .. _`Artur Balabanov`: https://github.com/arturbalabanov
 .. _`Kenny Do`: https://github.com/kennydo
+.. _`Max Rothman`: https://github.com/maxrothman
diff --git a/pgcli/main.py b/pgcli/main.py
index 4219930f..87b0ec51 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -42,7 +42,7 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from pygments.lexers.sql import PostgresLexer
 from pygments.token import Token
 
-from pgspecial.main import (PGSpecial, NO_QUERY, PAGER_OFF)
+from pgspecial.main import (PGSpecial, NO_QUERY, PAGER_OFF, PAGER_LONG_OUTPUT)
 import pgspecial as special
 try:
     import keyring
@@ -77,6 +77,9 @@ from collections import namedtuple
 
 from textwrap import dedent
 
+# Ref: https://stackoverflow.com/questions/30425105/filter-special-chars-such-as-color-codes-from-shell-output
+COLOR_CODE_REGEX = re.compile(r'\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))')
+
 # Query tuples are used for maintaining history
 MetaQuery = namedtuple(
     'Query',
@@ -929,6 +932,15 @@ class PGCli(object):
     def echo_via_pager(self, text, color=None):
         if self.pgspecial.pager_config == PAGER_OFF or self.watch_command:
             click.echo(text, color=color)
+        elif self.pgspecial.pager_config == PAGER_LONG_OUTPUT:
+            lines = text.split('\n')
+
+            # The last 4 lines are reserved for the pgcli menu and padding
+            if len(lines) >= self.cli.output.get_size().rows - 4 \
+                    or any(len(COLOR_CODE_REGEX.sub('', l)) > self.cli.output.get_size().columns for l in lines):
+                click.echo_via_pager(text, color=color)
+            else:
+                click.echo(text, color=color)
         else:
             click.echo_via_pager(text, color)
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index d01f54fa..575a7741 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,4 +1,5 @@
-pytest>=2.7.0
+# The maximum version requirement can be removed once Python 3.4 goes EOL
+pytest>=2.7.0,<=3.0.7
 mock>=1.0.1
 tox>=1.9.2
 behave>=1.2.4
diff --git a/tests/features/steps/wrappers.py b/tests/features/steps/wrappers.py
index de49cede..9a7f89da 100644
--- a/tests/features/steps/wrappers.py
+++ b/tests/features/steps/wrappers.py
@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 
 import re
 import pexpect
+from pgcli.main import COLOR_CODE_REGEX
 
 
 def expect_exact(context, expected, timeout):
@@ -10,7 +11,7 @@ def expect_exact(context, expected, timeout):
         context.cli.expect_exact(expected, timeout=timeout)
     except:
         # Strip color codes out of the output.
-        actual = re.sub(r'\x1b\[([0-9A-Za-z;?])+[m|K]?', '', context.cli.before)
+        actual = COLOR_CODE_REGEX.sub('', context.cli.before)
         raise Exception('Expected:\n---\n{0!r}\n---\n\nActual:\n---\n{1!r}\n---'.format(
             expected,
             actual))
diff --git a/tests/test_main.py b/tests/test_main.py
index 7d6d773e..905bd2e5 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -11,7 +11,7 @@ except ImportError:
     setproctitle = None
 
 from pgcli.main import (
-    obfuscate_process_password, format_output, PGCli, OutputSettings
+    obfuscate_process_password, format_output, PGCli, OutputSettings, COLOR_CODE_REGEX
 )
 from pgspecial.main import (PAGER_OFF, PAGER_LONG_OUTPUT, PAGER_ALWAYS)
 from tests.utils import dbtest, run
@@ -148,6 +148,100 @@ def test_format_output_auto_expand():
     assert '\n'.join(expanded_results) == '\n'.join(expanded)
 
 
+termsize = namedtuple('termsize', ['rows', 'columns'])
+test_line = '-' * 10
+test_data = [
+    (10, 10, '\n'.join([test_line] * 7)),
+    (10, 10, '\n'.join([test_line] * 6)),
+    (10, 10, '\n'.join([test_line] * 5)),
+    (10, 10, '-' * 11),
+    (10, 10, '-' * 10),
+    (10, 10, '-' * 9),
+]
+
+# 4 lines are reserved at the bottom of the terminal for pgcli's prompt
+use_pager_when_on = [True,
+                     True,
+                     False,
+                     True,
+                     False,
+                     False]
+
+# Can be replaced with pytest.param once we can upgrade pytest after Python 3.4 goes EOL
+test_ids = ["Output longer than terminal height",
+            "Output equal to terminal height",
+            "Output shorter than terminal height",
+            "Output longer than terminal width",
+            "Output equal to terminal width",
+            "Output shorter than terminal width"]
+
+
+@pytest.fixture
+def pset_pager_mocks():
+    cli = PGCli()
+    cli.watch_command = None
+    with mock.patch('pgcli.main.click.echo') as mock_echo, \
+            mock.patch('pgcli.main.click.echo_via_pager') as mock_echo_via_pager, \
+            mock.patch.object(cli, 'cli') as mock_cli:
+
+        yield cli, mock_echo, mock_echo_via_pager, mock_cli
+
+
+@pytest.mark.parametrize('term_height,term_width,text', test_data, ids=test_ids)
+def test_pset_pager_off(term_height, term_width, text, pset_pager_mocks):
+    cli, mock_echo, mock_echo_via_pager, mock_cli = pset_pager_mocks
+    mock_cli.output.get_size.return_value = termsize(
+        rows=term_height, columns=term_width)
+
+    with mock.patch.object(cli.pgspecial, 'pager_config', PAGER_OFF):
+        cli.echo_via_pager(text)
+
+    mock_echo.assert_called()
+    mock_echo_via_pager.assert_not_called()
+
+
+@pytest.mark.parametrize('term_height,term_width,text', test_data, ids=test_ids)
+def test_pset_pager_always(term_height, term_width, text, pset_pager_mocks):
+    cli, mock_echo, mock_echo_via_pager, mock_cli = pset_pager_mocks
+    mock_cli.output.get_size.return_value = termsize(
+        rows=term_height, columns=term_width)
+
+    with mock.patch.object(cli.pgspecial, 'pager_config', PAGER_ALWAYS):
+        cli.echo_via_pager(text)
+
+    mock_echo.assert_not_called()
+    mock_echo_via_pager.assert_called()
+
+
+pager_on_test_data = [l + (r,) for l, r in zip(test_data, use_pager_when_on)]
+
+
+@pytest.mark.parametrize('term_height,term_width,text,use_pager', pager_on_test_data, ids=test_ids)
+def test_pset_pager_on(term_height, term_width, text, use_pager, pset_pager_mocks):
+    cli, mock_echo, mock_echo_via_pager, mock_cli = pset_pager_mocks
+    mock_cli.output.get_size.return_value = termsize(
+        rows=term_height, columns=term_width)
+
+    with mock.patch.object(cli.pgspecial, 'pager_config', PAGER_LONG_OUTPUT):
+        cli.echo_via_pager(text)
+
+    if use_pager:
+        mock_echo.assert_not_called()
+        mock_echo_via_pager.assert_called()
+    else:
+        mock_echo_via_pager.assert_not_called()
+        mock_echo.assert_called()
+
+
+@pytest.mark.parametrize('text,expected_length', [
+    (u"22200K .......\u001b[0m\u001b[91m... .......... ...\u001b[0m\u001b[91m.\u001b[0m\u001b[91m...... .........\u001b[0m\u001b[91m.\u001b[0m\u001b[91m \u001b[0m\u001b[91m.\u001b[0m\u001b[91m.\u001b[0m\u001b[91m.\u001b[0m\u001b[91m.\u001b[0m\u001b[91m...... 50% 28.6K 12m55s", 78),
+    (u"=\u001b[m=", 2),
+    (u"-\u001b]23\u0007-", 2),
+])
+def test_color_pattern(text, expected_length, pset_pager_mocks):
+    cli = pset_pager_mocks[0]
+    assert len(COLOR_CODE_REGEX.sub('', text)) == expected_length
+
 @dbtest
 def test_i_works(tmpdir, executor):
     sqlfile = tmpdir.join("test.sql")