Add tests for .csv()

Author: catherinedevlin

src/sql/magic.py

@@ -116,17 +116,25 @@ def execute(self, line, cell='', local_ns={}):
 
     legal_sql_identifier = re.compile(r'^[A-Za-z0-9#_$]+')
     def _persist_dataframe(self, raw, conn, user_ns):
+        """Implements PERSIST, which writes a DataFrame to the RDBMS"""
         if not DataFrame:
             raise ImportError("Must `pip install pandas` to use DataFrames")
+
+        # Parse input to get name of DataFrame
         pieces = raw.split()
         if len(pieces) != 2:
             raise SyntaxError("Format: %sql [connection] persist <DataFrameName>")
         frame_name = pieces[1].strip(';')
+
+       # Get the DataFrame from the user namespace
         frame = eval(frame_name, user_ns)
         if not isinstance(frame, DataFrame) and not isinstance(frame, Series):
             raise TypeError('%s is not a Pandas DataFrame or Series' % frame_name)
+
+       # Make a suitable name for the resulting database table
         table_name = frame_name.lower()
         table_name = self.legal_sql_identifier.search(table_name).group(0)
+
         frame.to_sql(table_name, conn.session.engine)
         return 'Persisted %s' % table_name
 

src/sql/parse.py

@@ -4,9 +4,12 @@
 
 
 def parse(cell, config):
+    """Separate input into (connection info, SQL statement)"""
+
     parts = [part.strip() for part in cell.split(None, 1)]
     if not parts:
         return {'connection': '', 'sql': ''}
+
     if parts[0].startswith('[') and parts[0].endswith(']'):
         section = parts[0].lstrip('[').rstrip(']')
         parser = CP.ConfigParser()

src/tests/test_magic.py

@@ -1,7 +1,9 @@
 from nose import with_setup
 from sql.magic import SqlMagic
 from textwrap import dedent
+import os.path
 import re
+import tempfile
 
 ip = get_ipython()
 
@@ -132,3 +134,26 @@ def test_autopandas():
     dframe = ip.run_cell("%sql SELECT * FROM test;")
     assert dframe.success
     assert dframe.result.name[0] == 'foo'
+
+@with_setup(_setup, _teardown)
+def test_csv():
+    ip.run_line_magic('config',  "SqlMagic.autopandas = False")  # uh-oh
+    result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM test;")
+    result = result.csv()
+    for row in result.splitlines():
+        assert row.count(',') == 1
+    assert len(result.splitlines()) == 3
+
+@with_setup(_setup, _teardown)
+def test_csv_to_file():
+    ip.run_line_magic('config',  "SqlMagic.autopandas = False")  # uh-oh
+    result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM test;")
+    with tempfile.TemporaryDirectory() as tempdir:
+        fname = os.path.join(tempdir, 'test.csv')
+        output = result.csv(fname)
+        assert os.path.exists(output.file_path)
+        with open(output.file_path) as csvfile:
+            content = csvfile.read()
+            for row in content.splitlines():
+                assert row.count(',') == 1
+            assert len(content.splitlines()) == 3