Merge pull request #70 from xiaochuanyu/master

Added a way to store result in specified variable without printing it.

Author: catherinedevlin

.gitignore

@@ -33,3 +33,6 @@ nosetests.xml
 .mr.developer.cfg
 .project
 .pydevproject
+
+# Pycharm
+/.idea

src/sql/magic.py

@@ -78,9 +78,10 @@ def execute(self, line, cell='', local_ns={}):
         user_ns.update(local_ns)
 
         parsed = sql.parse.parse('%s\n%s' % (line, cell), self)
+        flags = parsed['flags']
         conn = sql.connection.Connection.get(parsed['connection'])
-        first_word = parsed['sql'].split(None, 1)[:1]
-        if first_word and first_word[0].lower() == 'persist':
+
+        if flags['persist']:
             return self._persist_dataframe(parsed['sql'], conn, user_ns)
 
         try:
@@ -104,6 +105,13 @@ def execute(self, line, cell='', local_ns={}):
 
                 return None
             else:
+
+                if flags['result_var']:
+                    result_var = flags['result_var']
+                    print("Returning data to local variable {}".format(result_var))
+                    self.shell.user_ns.update({result_var: result})
+                    return None
+
                 #Return results into the default ipython _ variable
                 return result
 
@@ -118,10 +126,9 @@ def execute(self, line, cell='', local_ns={}):
     def _persist_dataframe(self, raw, conn, user_ns):
         if not DataFrame:
             raise ImportError("Must `pip install pandas` to use DataFrames")
-        pieces = raw.split()
-        if len(pieces) != 2:
+        if not raw:
             raise SyntaxError("Format: %sql [connection] persist <DataFrameName>")
-        frame_name = pieces[1].strip(';')
+        frame_name = raw.strip(';')
         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)

src/sql/parse.py

@@ -24,5 +24,27 @@ def parse(cell, config):
     else:
         connection = ''
         sql = cell
+    flags, sql = parse_sql_flags(sql.strip())
     return {'connection': connection.strip(),
-            'sql': sql.strip()}
+            'sql': sql,
+            'flags': flags}
+
+
+def parse_sql_flags(sql):
+    words = sql.split()
+    flags = {
+        'persist': False,
+        'result_var': None
+    }
+    if not words:
+        return (flags, "")
+    num_words = len(words)
+    trimmed_sql = sql
+    if words[0].lower() == 'persist':
+        print("Persist parsed to True")
+        flags['persist'] = True
+        trimmed_sql =  " ".join(words[1:])
+    elif num_words >= 2 and words[1] == '<<':
+        flags['result_var'] = words[0]
+        trimmed_sql = " ".join(words[2:])
+    return (flags, trimmed_sql.strip())

src/tests/test_magic.py

@@ -55,6 +55,16 @@ def test_multi_sql():
         """)
     assert 'Shakespeare' in str(result) and 'Brecht' in str(result)
 
+@with_setup(_setup_writer, _teardown_writer)
+def test_result_var():
+    ip.run_cell_magic('sql', '', """
+        sqlite://
+        x <<
+        SELECT last_name FROM writer;
+        """)
+    result = ip.user_global_ns['x']
+    assert 'Shakespeare' in str(result) and 'Brecht' in str(result)
+
 @with_setup(_setup_writer, _teardown_writer)
 def test_access_results_by_keys():
     assert ip.run_line_magic('sql', "sqlite:// SELECT * FROM writer;")['William'] == (u'William', u'Shakespeare', 1616)
@@ -76,7 +86,6 @@ def test_autolimit():
     result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM test;")
     assert len(result) == 1
 
-
 @with_setup(_setup, _teardown)
 def test_persist():
     ip.run_cell("results = %sql SELECT * FROM test;")

src/tests/test_parse.py

@@ -6,24 +6,28 @@
     from IPython.config.configurable import Configurable
 
 empty_config = Configurable()
-
+default_flags = {'persist': False, 'result_var': None}
 def test_parse_no_sql():
     assert parse("will:longliveliz@localhost/shakes", empty_config) == \
            {'connection': "will:longliveliz@localhost/shakes",
-            'sql': ''}
+            'sql': '',
+            'flags': default_flags}
 
 def test_parse_with_sql():
     assert parse("postgresql://will:longliveliz@localhost/shakes SELECT * FROM work", 
                  empty_config) == \
            {'connection': "postgresql://will:longliveliz@localhost/shakes",
-            'sql': 'SELECT * FROM work'}    
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags}    
 
 def test_parse_sql_only():
     assert parse("SELECT * FROM work", empty_config) == \
            {'connection': "",
-            'sql': 'SELECT * FROM work'} 
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags} 
 
 def test_parse_postgresql_socket_connection():
     assert parse("postgresql:///shakes SELECT * FROM work", empty_config) == \
            {'connection': "postgresql:///shakes",
-            'sql': 'SELECT * FROM work'}            
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags}