Adjust displaylimit even of result sets already fetched

Author: catherinedevlin

HACKING.txt

@@ -1,17 +1,23 @@
 Development setup
 =================
 
-Running nose tests with IPython is tricky, so there's a 
+Running nose tests with IPython is tricky, so there's a
 run_tests.sh script for it.
 
+To temporarily insert breakpoints for debugging: `from nose.tools import set_trace; set_trace()`
+
+Tests have requirements not installed by setup.py:
+
+- nose
+- pandas
+
 Release HOWTO
 =============
 
-To make a release, 
+To make a release,
 
   1) Update release date/version in NEWS.txt and setup.py
   2) Run 'python setup.py sdist'
   3) Test the generated source distribution in dist/
   4) Upload to PyPI: 'python setup.py sdist register upload'
   5) Increase version in setup.py (for next release)
-

src/sql/magic.py

@@ -26,10 +26,10 @@ class SqlMagic(Magics, Configurable):
 
     Provides the %%sql magic."""
 
-    autolimit = Int(0, config=True, help="Automatically limit the size of the returned result sets")
+    autolimit = Int(0, config=True, allow_none=True, help="Automatically limit the size of the returned result sets")
     style = Unicode('DEFAULT', config=True, help="Set the table printing style to any of prettytable's defined styles (currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM)")
     short_errors = Bool(True, config=True, help="Don't display the full traceback on SQL Programming Error")
-    displaylimit = Int(0, config=True, help="Automatically limit the number of rows displayed (full result set is still stored)")
+    displaylimit = Int(None, config=True, allow_none=True, help="Automatically limit the number of rows displayed (full result set is still stored)")
     autopandas = Bool(False, config=True, help="Return Pandas DataFrames instead of regular result sets")
     column_local_vars = Bool(False, config=True, help="Return data into local variables from column names")
     feedback = Bool(True, config=True, help="Print number of rows affected by DML")
@@ -157,7 +157,7 @@ def load_ipython_extension(ip):
 
     # this fails in both Firefox and Chrome for OS X.
     # I get the error: TypeError: IPython.CodeCell.config_defaults is undefined
-    
+
     # js = "IPython.CodeCell.config_defaults.highlight_modes['magic_sql'] = {'reg':[/^%%sql/]};"
     # display_javascript(js, raw=True)
     ip.register_magics(SqlMagic)

src/sql/run.py

@@ -103,17 +103,16 @@ def __init__(self, sqlaproxy, sql, config):
             else:
                 list.__init__(self, sqlaproxy.fetchall())
             self.field_names = unduplicate_field_names(self.keys)
-            self.pretty = prettytable.PrettyTable(self.field_names)
-            if not config.autopandas:
-                for row in self[:config.displaylimit or None]:
-                    self.pretty.add_row(row)
-            self.pretty.set_style(self.style)
+            self.pretty = PrettyTable(self.field_names, style=self.style)
+            # self.pretty.set_style(self.style)
         else:
             list.__init__(self, [])
             self.pretty = None
+
     def _repr_html_(self):
         _cell_with_spaces_pattern = re.compile(r'(<td>)( {2,})')
         if self.pretty:
+            self.pretty.add_rows(self)
             result = self.pretty.get_html_string()
             result = _cell_with_spaces_pattern.sub(_nonbreaking_spaces, result)
             if self.config.displaylimit and len(self) > self.config.displaylimit:
@@ -122,8 +121,11 @@ def _repr_html_(self):
             return result
         else:
             return None
+
     def __str__(self, *arg, **kwarg):
+        self.pretty.add_rows(self)
         return str(self.pretty or '')
+
     def __getitem__(self, key):
         """
         Access by integer (row position within result set)
@@ -247,6 +249,7 @@ def csv(self, filename=None, **format_params):
            Any other parameters will be passed on to csv.writer."""
         if not self.pretty:
             return None # no results
+        self.pretty.add_rows(self)
         if filename:
             encoding = format_params.get('encoding', 'utf-8')
             if six.PY2:
@@ -316,3 +319,25 @@ def run(conn, sql, config, user_namespace):
         #returning only last result, intentionally
     else:
         return 'Connected: %s' % conn.name
+
+
+class PrettyTable(prettytable.PrettyTable):
+
+    def __init__(self, *args, **kwargs):
+        self.row_count = 0
+        self.displaylimit = None
+        return super(PrettyTable, self).__init__(*args,  **kwargs)
+
+    def add_rows(self, data):
+        if self.row_count and (data.config.displaylimit == self.displaylimit):
+            return  # correct number of rows already present
+        self.clear_rows()
+        self.displaylimit = data.config.displaylimit
+        if self.displaylimit == 0:
+            self.displaylimit = None  # TODO: remove this to make 0 really 0
+        if self.displaylimit in (None, 0):
+            self.row_count = len(data)
+        else:
+            self.row_count = min(len(data), self.displaylimit)
+        for row in data[:self.displaylimit]:
+            self.add_row(row)

src/tests/test_magic.py

@@ -39,7 +39,7 @@ def test_print():
 def test_plain_style():
     ip.run_line_magic('config',  "SqlMagic.style = 'PLAIN_COLUMNS'")
     result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM test;")
-    assert re.search(r'1\s+foo', str(result))
+    assert re.search(r'1\s+\|\s+foo', str(result))
 
 
 def _setup_writer():
@@ -120,23 +120,24 @@ def test_persist_frame_at_its_creation():
     persisted = ip.run_line_magic('sql', 'SELECT * FROM results')
     assert 'Shakespeare' in str(persisted)
 
-# TODO: support 
+# TODO: support
 # @with_setup(_setup_writer, _teardown_writer)
 # def test_persist_with_connection_info():
 #     ip.run_cell("results = %sql SELECT * FROM writer;")
 #     ip.run_line_magic('sql', 'sqlite:// PERSIST results.DataFrame()')
 #     persisted = ip.run_line_magic('sql', 'SELECT * FROM results')
 #     assert 'Shakespeare' in str(persisted)
 
-@with_setup(_setup_writer, _teardown_writer)
 def test_displaylimit():
-    ip.run_line_magic('config',  "SqlMagic.autolimit = 0")
-    ip.run_line_magic('config',  "SqlMagic.displaylimit = 0")
-    result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM writer;")
-    assert result._repr_html_().count("<tr>") == 3
+    ip.run_line_magic('config',  "SqlMagic.autolimit = None")
+    ip.run_line_magic('config',  "SqlMagic.displaylimit = None")
+    result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM (VALUES ('apple'), ('banana'), ('cherry')) AS Result ORDER BY 1;")
+    assert 'apple' in result._repr_html_()
+    assert 'banana' in result._repr_html_()
+    assert 'cherry' in result._repr_html_()
     ip.run_line_magic('config',  "SqlMagic.displaylimit = 1")
-    result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM writer;")
-    assert result._repr_html_().count("<tr>") == 2
+    assert 'apple' in result._repr_html_()
+    assert 'cherry' not in result._repr_html_()
 
 @with_setup(_setup_writer, _teardown_writer)
 def test_column_local_vars():
@@ -210,4 +211,3 @@ def test_dicts():
         assert 'first_name' in row
         assert 'last_name' in row
         assert 'year_of_death' in row
-