Merge pull request #559 from python-cmd2/self

Remove self from pystate if locals_in_py is False

2 files changed, 20 insertions, 1 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 03c71a91..e04138fa 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -2982,6 +2982,8 @@ class Cmd(cmd.Cmd):
 
             if self.locals_in_py:
                 self.pystate['self'] = self
+            elif 'self' in self.pystate:
+                del self.pystate['self']
 
             localvars = self.pystate
             from code import InteractiveConsole
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
index a382a940..3d57a105 100644
--- a/tests/test_cmd2.py
+++ b/tests/test_cmd2.py
@@ -216,18 +216,35 @@ def test_base_shell(base_app, monkeypatch):
     assert m.called
 
 def test_base_py(base_app, capsys):
+    # Create a variable and make sure we can see it
     run_cmd(base_app, 'py qqq=3')
     out, err = capsys.readouterr()
     assert out == ''
-
     run_cmd(base_app, 'py print(qqq)')
     out, err = capsys.readouterr()
     assert out.rstrip() == '3'
 
+    # Add a more complex statement
     run_cmd(base_app, 'py print("spaces" + " in this " + "command")')
     out, err = capsys.readouterr()
     assert out.rstrip() == 'spaces in this command'
 
+    # Set locals_in_py to True and make sure we see self
+    out = run_cmd(base_app, 'set locals_in_py True')
+    assert 'now: True' in out
+
+    run_cmd(base_app, 'py print(self)')
+    out, err = capsys.readouterr()
+    assert 'cmd2.cmd2.Cmd object' in out
+
+    # Set locals_in_py to False and make sure we can't see self
+    out = run_cmd(base_app, 'set locals_in_py False')
+    assert 'now: False' in out
+
+    run_cmd(base_app, 'py print(self)')
+    out, err = capsys.readouterr()
+    assert "name 'self' is not defined" in err
+
 
 @pytest.mark.skipif(sys.platform == 'win32',
                     reason="Unit test doesn't work on win32, but feature does")