output via generator

changelog.md

@@ -12,6 +12,7 @@ Internal Changes:
 -----------------
 
 * Use less memory when formatting results for display (Thanks: [Dick Marinus]).
+* Preliminary work for a future change in outputting results that uses less memory (Thanks: [Dick Marinus]).
 
 1.12.0:
 =======

mycli/main.py

@@ -44,6 +44,8 @@
 from .lexer import MyCliLexer
 from .__init__ import __version__
 
+import itertools
+
 click.disable_unicode_literals_warning = True
 
 try:
@@ -543,7 +545,7 @@ def one_iteration(document=None):
                         if result_count > 0:
                             self.echo('')
                         try:
-                            self.output('\n'.join(formatted), status)
+                            self.output(formatted, status)
                         except KeyboardInterrupt:
                             pass
                         if special.is_timing_enabled():
@@ -678,8 +680,9 @@ def echo(self, s, **kwargs):
         self.log_output(s)
         click.secho(s, **kwargs)
 
-    def output_fits_on_screen(self, output, status=None):
-        """Check if the given output fits on the screen."""
+    def get_output_margin(self, status=None):
+        """Get the output margin (number of rows for the prompt, footer and
+        timing message."""
         size = self.cli.output.get_size()
 
         margin = self.get_reserved_space() + self.get_prompt(self.prompt_format).count('\n') + 1
@@ -688,11 +691,8 @@ def output_fits_on_screen(self, output, status=None):
         if status:
             margin += 1 + status.count('\n')
 
-        for i, line in enumerate(output.splitlines(), 1):
-            if len(line) > size.columns or i > (size.rows - margin):
-                return False
+        return margin
 
-        return True
 
     def output(self, output, status=None):
         """Output text to stdout or a pager command.
@@ -705,15 +705,42 @@ def output(self, output, status=None):
 
         """
         if output:
-            self.log_output(output)
-            special.write_tee(output)
-            special.write_once(output)
+            size = self.cli.output.get_size()
+
+            margin = self.get_output_margin(status)
+
+            fits = True
+            buf = []
+            output_via_pager = self.explicit_pager and special.is_pager_enabled()
+            for i, line in enumerate(output, 1):
+                self.log_output(line)
+                special.write_tee(line)
+                special.write_once(line)
+
+                if fits or output_via_pager:
+                    # buffering
+                    buf.append(line)
+                    if len(line) > size.columns or i > (size.rows - margin):
+                        fits = False
+                        if not self.explicit_pager and special.is_pager_enabled():
+                            # doesn't fit, use pager
+                            output_via_pager = True
+
+                        if not output_via_pager:
+                            # doesn't fit, flush buffer
+                            for line in buf:
+                                click.secho(line)
+                            buf = []
+                else:
+                    click.secho(line)
 
-            if (self.explicit_pager or
-                    (special.is_pager_enabled() and not self.output_fits_on_screen(output, status))):
-                click.echo_via_pager(output)
-            else:
-                click.secho(output)
+            if buf:
+                if output_via_pager:
+                    # sadly click.echo_via_pager doesn't accept generators
+                    click.echo_via_pager("\n".join(buf))
+                else:
+                    for line in buf:
+                        click.secho(line)
 
         if status:
             self.log_output(status)
@@ -812,7 +839,7 @@ def format_output(self, title, cur, headers, expanded=False,
         }
 
         if title:  # Only print the title if it's not None.
-            output.append(title)
+            output = itertools.chain(output, [title])
 
         if cur:
             column_types = None
@@ -829,14 +856,23 @@ def sanitize(col):
                 cur, headers, format_name='vertical' if expanded else None,
                 column_types=column_types,
                 **output_kwargs)
-            first_line = formatted[:formatted.find('\n')]
+
+            if isinstance(formatted, (text_type)):
+                formatted = formatted.splitlines()
+            formatted = iter(formatted)
+
+            first_line = next(formatted)
+            formatted = itertools.chain([first_line], formatted)
 
             if (not expanded and max_width and headers and cur and
                     len(first_line) > max_width):
                 formatted = self.formatter.format_output(
                     cur, headers, format_name='vertical', column_types=column_types, **output_kwargs)
+                if isinstance(formatted, (text_type)):
+                    formatted = iter(formatted.splitlines())
+
+            output = itertools.chain(output, formatted)
 
-            output.append(formatted)
 
         return output
 
@@ -988,7 +1024,6 @@ def cli(database, user, host, port, socket, password, dbname,
 
             if csv:
                 mycli.formatter.format_name = 'csv'
-                new_line = False
             elif not table:
                 mycli.formatter.format_name = 'tsv'
 

test/features/steps/auto_vertical.py

@@ -40,7 +40,7 @@ def step_see_large_results(context):
     rows = ['{n:3}| {n}'.format(n=str(n)) for n in range(1, 50)]
     expected = ('***************************[ 1. row ]'
                 '***************************\r\n' +
-                '{}\r\n\r\n'.format('\r\n'.join(rows)))
+                '{}\r\n'.format('\r\n'.join(rows)))
 
     wrappers.expect_pager(context, expected, timeout=5)
     wrappers.expect_exact(context, '1 row in set', timeout=2)

test/test_main.py

@@ -10,6 +10,7 @@
 from utils import USER, HOST, PORT, PASSWORD, dbtest, run
 
 from textwrap import dedent
+from collections import namedtuple
 
 try:
     text_type = basestring
@@ -64,10 +65,10 @@ def test_execute_arg_with_csv(executor):
     sql = 'select * from test;'
     runner = CliRunner()
     result = runner.invoke(cli, args=CLI_ARGS + ['-e', sql] + ['--csv'])
-    expected = 'a\nabc\n\n'
+    expected = 'a\nabc\n'
 
     assert result.exit_code == 0
-    assert expected in result.output
+    assert expected in "".join(result.output)
 
 
 @dbtest
@@ -84,7 +85,7 @@ def test_batch_mode(executor):
     result = runner.invoke(cli, args=CLI_ARGS, input=sql)
 
     assert result.exit_code == 0
-    assert 'count(*)\n3\n\na\nabc\n' in result.output
+    assert 'count(*)\n3\na\nabc\n' in "".join(result.output)
 
 
 @dbtest
@@ -129,7 +130,7 @@ def test_batch_mode_csv(executor):
     expected = 'a,b\nabc,def\nghi,jkl\n'
 
     assert result.exit_code == 0
-    assert expected in result.output
+    assert expected in "".join(result.output)
 
 
 @dbtest
@@ -201,6 +202,98 @@ def test_command_descriptions_end_with_periods():
         assert command[3].endswith('.')
 
 
+def output(monkeypatch, terminal_size, testdata, explicit_pager, expect_pager):
+    global clickoutput
+    clickoutput = ""
+    m = MyCli()
+
+    class TestOutput():
+        def get_size(self):
+            size = namedtuple('Size', 'rows columns')
+            size.columns, size.rows = terminal_size
+            return size
+
+    class TestExecute():
+        host = 'test'
+        user = 'test'
+        dbname = 'test'
+        port = 0
+
+        def server_type(self):
+            return ['test']
+
+    class CommandLineInterface():
+        output = TestOutput()
+
+    m.cli = CommandLineInterface()
+    m.sqlexecute = TestExecute()
+    m.explicit_pager = explicit_pager
+
+    def echo_via_pager(s):
+        assert expect_pager
+        global clickoutput
+        clickoutput += s
+
+    def secho(s):
+        assert not expect_pager
+        global clickoutput
+        clickoutput += s + "\n"
+
+    monkeypatch.setattr(click, 'echo_via_pager', echo_via_pager)
+    monkeypatch.setattr(click, 'secho', secho)
+    m.output(testdata)
+    if clickoutput.endswith("\n"):
+        clickoutput = clickoutput[:-1]
+    assert clickoutput == "\n".join(testdata)
+
+
+def test_conditional_pager(monkeypatch):
+    testdata = "Lorem ipsum dolor sit amet consectetur adipiscing elit sed do".split(
+        " ")
+    # User didn't set pager, output doesn't fit screen -> pager
+    output(
+        monkeypatch,
+        terminal_size=(5, 10),
+        testdata=testdata,
+        explicit_pager=False,
+        expect_pager=True
+    )
+    # User didn't set pager, output fits screen -> no pager
+    output(
+        monkeypatch,
+        terminal_size=(20, 20),
+        testdata=testdata,
+        explicit_pager=False,
+        expect_pager=False
+    )
+    # User manually configured pager, output doesn't fit screen -> pager
+    output(
+        monkeypatch,
+        terminal_size=(5, 10),
+        testdata=testdata,
+        explicit_pager=True,
+        expect_pager=True
+    )
+    # User manually configured pager, output fit screen -> pager
+    output(
+        monkeypatch,
+        terminal_size=(20, 20),
+        testdata=testdata,
+        explicit_pager=True,
+        expect_pager=True
+    )
+
+    SPECIAL_COMMANDS['nopager'].handler()
+    output(
+        monkeypatch,
+        terminal_size=(5, 10),
+        testdata=testdata,
+        explicit_pager=False,
+        expect_pager=False
+    )
+    SPECIAL_COMMANDS['pager'].handler('')
+
+
 def test_reserved_space_is_integer():
     """Make sure that reserved space is returned as an integer."""
     def stub_terminal_size():