diff --git a/src/sql/magic.py b/src/sql/magic.py index e073b29da..25d04289a 100644 --- a/src/sql/magic.py +++ b/src/sql/magic.py @@ -586,23 +586,23 @@ def interactive_execute_wrapper(**kwargs): and not isinstance(result, str) and self.column_local_vars ): - # Instead of returning values, set variables directly in the - # users namespace. Variable names given by column names + # Set variables directly in the users namespace. + # Variable names given by column names if self.autopandas or self.autopolars: keys = result.keys() + self.shell.user_ns.update(result) + result = None else: keys = result.keys - result = result.dict() + self.shell.user_ns.update(result.dict()) if self.feedback: display.message( "Returning data to local variables [{}]".format(", ".join(keys)) ) - self.shell.user_ns.update(result) - - return None + return result else: if command.result_var: self.shell.user_ns.update({command.result_var: result}) diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py index 0b0c0af93..55d73e59a 100644 --- a/src/tests/test_magic.py +++ b/src/tests/test_magic.py @@ -737,7 +737,7 @@ def test_displaylimit_with_count_statement(ip, load_penguin, config_value): def test_column_local_vars(ip): ip.run_line_magic("config", "SqlMagic.column_local_vars = True") result = runsql(ip, "SELECT * FROM author;") - assert result is None + assert result is not None assert "William" in ip.user_global_ns["first_name"] assert "Shakespeare" in ip.user_global_ns["last_name"] assert len(ip.user_global_ns["first_name"]) == 2