Issue *18081, #18242: Change Idle warnings capture in PyShell and run to stop

replacing warnings.formatwarnings and to reverse replacement of
warnings.showwarnings when import is complete and when main function exits.
Add test_warning.py. Vinay Sajip provided capture_warnings function.

3 files changed, 162 insertions, 49 deletions

diff --git a/Lib/idlelib/PyShell.py b/Lib/idlelib/PyShell.py
index 7f2b775..b894462 100644
--- a/Lib/idlelib/PyShell.py
+++ b/Lib/idlelib/PyShell.py
@@ -45,35 +45,55 @@ PORT = 0  # someday pass in host, port for remote debug capability
 # internal warnings to the console.  ScriptBinding.check_syntax() will
 # temporarily redirect the stream to the shell window to display warnings when
 # checking user's code.
-global warning_stream
-warning_stream = sys.__stderr__
-try:
-    import warnings
-except ImportError:
-    pass
-else:
-    def idle_showwarning(message, category, filename, lineno,
-                         file=None, line=None):
-        if file is None:
-            file = warning_stream
-        try:
-            file.write(warnings.formatwarning(message, category, filename,
-                                              lineno, line=line))
-        except OSError:
-            pass  ## file (probably __stderr__) is invalid, warning dropped.
-    warnings.showwarning = idle_showwarning
-    def idle_formatwarning(message, category, filename, lineno, line=None):
-        """Format warnings the IDLE way"""
-        s = "\nWarning (from warnings module):\n"
-        s += '  File \"%s\", line %s\n' % (filename, lineno)
-        if line is None:
-            line = linecache.getline(filename, lineno)
-        line = line.strip()
-        if line:
-            s += "    %s\n" % line
-        s += "%s: %s\n>>> " % (category.__name__, message)
-        return s
-    warnings.formatwarning = idle_formatwarning
+warning_stream = sys.__stderr__  # None, at least on Windows, if no console.
+import warnings
+
+def idle_formatwarning(message, category, filename, lineno, line=None):
+    """Format warnings the IDLE way."""
+
+    s = "\nWarning (from warnings module):\n"
+    s += '  File \"%s\", line %s\n' % (filename, lineno)
+    if line is None:
+        line = linecache.getline(filename, lineno)
+    line = line.strip()
+    if line:
+        s += "    %s\n" % line
+    s += "%s: %s\n" % (category.__name__, message)
+    return s
+
+def idle_showwarning(
+        message, category, filename, lineno, file=None, line=None):
+    """Show Idle-format warning (after replacing warnings.showwarning).
+
+    The differences are the formatter called, the file=None replacement,
+    which can be None, the capture of the consequence AttributeError,
+    and the output of a hard-coded prompt.
+    """
+    if file is None:
+        file = warning_stream
+    try:
+        file.write(idle_formatwarning(
+                message, category, filename, lineno, line=line))
+        file.write(">>> ")
+    except (AttributeError, OSError):
+        pass  # if file (probably __stderr__) is invalid, skip warning.
+
+_warnings_showwarning = None
+
+def capture_warnings(capture):
+    "Replace warning.showwarning with idle_showwarning, or reverse."
+
+    global _warnings_showwarning
+    if capture:
+        if _warnings_showwarning is None:
+            _warnings_showwarning = warnings.showwarning
+            warnings.showwarning = idle_showwarning
+    else:
+        if _warnings_showwarning is not None:
+            warnings.showwarning = _warnings_showwarning
+            _warnings_showwarning = None
+
+capture_warnings(True)
 
 def extended_linecache_checkcache(filename=None,
                                   orig_checkcache=linecache.checkcache):
@@ -1421,6 +1441,7 @@ echo "import sys; print(sys.argv)" | idle - "foobar"
 def main():
     global flist, root, use_subprocess
 
+    capture_warnings(True)
     use_subprocess = True
     enable_shell = False
     enable_edit = False
@@ -1553,7 +1574,10 @@ def main():
     while flist.inversedict:  # keep IDLE running while files are open.
         root.mainloop()
     root.destroy()
+    capture_warnings(False)
 
 if __name__ == "__main__":
     sys.modules['PyShell'] = sys.modules['__main__']
     main()
+
+capture_warnings(False)  # Make sure turned off; see issue 18081
diff --git a/Lib/idlelib/idle_test/test_warning.py b/Lib/idlelib/idle_test/test_warning.py
new file mode 100644
index 0000000..18627dd
--- /dev/null
+++ b/Lib/idlelib/idle_test/test_warning.py
@@ -0,0 +1,73 @@
+'''Test warnings replacement in PyShell.py and run.py.
+
+This file could be expanded to include traceback overrides
+(in same two modules). If so, change name.
+Revise if output destination changes (http://bugs.python.org/issue18318).
+Make sure warnings module is left unaltered (http://bugs.python.org/issue18081).
+'''
+
+import unittest
+from test.support import captured_stderr
+
+import warnings
+# Try to capture default showwarning before Idle modules are imported.
+showwarning = warnings.showwarning
+# But if we run this file within idle, we are in the middle of the run.main loop
+# and default showwarnings has already been replaced.
+running_in_idle = 'idle' in showwarning.__name__
+
+from idlelib import run
+from idlelib import PyShell as shell
+
+# The following was generated from PyShell.idle_formatwarning
+# and checked as matching expectation.
+idlemsg = '''
+Warning (from warnings module):
+  File "test_warning.py", line 99
+    Line of code
+UserWarning: Test
+'''
+shellmsg = idlemsg + ">>> "
+
+class RunWarnTest(unittest.TestCase):
+
+    @unittest.skipIf(running_in_idle, "Does not work when run within Idle.")
+    def test_showwarnings(self):
+        self.assertIs(warnings.showwarning, showwarning)
+        run.capture_warnings(True)
+        self.assertIs(warnings.showwarning, run.idle_showwarning_subproc)
+        run.capture_warnings(False)
+        self.assertIs(warnings.showwarning, showwarning)
+
+    def test_run_show(self):
+        with captured_stderr() as f:
+            run.idle_showwarning_subproc(
+                    'Test', UserWarning, 'test_warning.py', 99, f, 'Line of code')
+            # The following uses .splitlines to erase line-ending differences
+            self.assertEqual(idlemsg.splitlines(), f.getvalue().splitlines())
+
+class ShellWarnTest(unittest.TestCase):
+
+    @unittest.skipIf(running_in_idle, "Does not work when run within Idle.")
+    def test_showwarnings(self):
+        self.assertIs(warnings.showwarning, showwarning)
+        shell.capture_warnings(True)
+        self.assertIs(warnings.showwarning, shell.idle_showwarning)
+        shell.capture_warnings(False)
+        self.assertIs(warnings.showwarning, showwarning)
+
+    def test_idle_formatter(self):
+        # Will fail if format changed without regenerating idlemsg
+        s = shell.idle_formatwarning(
+                'Test', UserWarning, 'test_warning.py', 99, 'Line of code')
+        self.assertEqual(idlemsg, s)
+
+    def test_shell_show(self):
+        with captured_stderr() as f:
+            shell.idle_showwarning(
+                    'Test', UserWarning, 'test_warning.py', 99, f, 'Line of code')
+            self.assertEqual(shellmsg.splitlines(), f.getvalue().splitlines())
+
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2, exit=False)
diff --git a/Lib/idlelib/run.py b/Lib/idlelib/run.py
index f5d260d..c1859b6 100644
--- a/Lib/idlelib/run.py
+++ b/Lib/idlelib/run.py
@@ -23,28 +23,39 @@ import __main__
 
 LOCALHOST = '127.0.0.1'
 
-try:
-    import warnings
-except ImportError:
-    pass
-else:
-    def idle_formatwarning_subproc(message, category, filename, lineno,
-                                   line=None):
-        """Format warnings the IDLE way"""
-        s = "\nWarning (from warnings module):\n"
-        s += '  File \"%s\", line %s\n' % (filename, lineno)
-        if line is None:
-            line = linecache.getline(filename, lineno)
-        line = line.strip()
-        if line:
-            s += "    %s\n" % line
-        s += "%s: %s\n" % (category.__name__, message)
-        return s
-    warnings.formatwarning = idle_formatwarning_subproc
+import warnings
 
+def idle_showwarning_subproc(
+        message, category, filename, lineno, file=None, line=None):
+    """Show Idle-format warning after replacing warnings.showwarning.
 
-tcl = tkinter.Tcl()
+    The only difference is the formatter called.
+    """
+    if file is None:
+        file = sys.stderr
+    try:
+        file.write(PyShell.idle_formatwarning(
+                message, category, filename, lineno, line))
+    except IOError:
+        pass # the file (probably stderr) is invalid - this warning gets lost.
 
+_warnings_showwarning = None
+
+def capture_warnings(capture):
+    "Replace warning.showwarning with idle_showwarning_subproc, or reverse."
+
+    global _warnings_showwarning
+    if capture:
+        if _warnings_showwarning is None:
+            _warnings_showwarning = warnings.showwarning
+            warnings.showwarning = idle_showwarning_subproc
+    else:
+        if _warnings_showwarning is not None:
+            warnings.showwarning = _warnings_showwarning
+            _warnings_showwarning = None
+
+capture_warnings(True)
+tcl = tkinter.Tcl()
 
 def handle_tk_events(tcl=tcl):
     """Process any tk events that are ready to be dispatched if tkinter
@@ -52,7 +63,6 @@ def handle_tk_events(tcl=tcl):
     loaded."""
     tcl.eval("update")
 
-
 # Thread shared globals: Establish a queue between a subthread (which handles
 # the socket) and the main thread (which runs user code), plus global
 # completion, exit and interruptable (the main thread) flags:
@@ -91,6 +101,8 @@ def main(del_exitfunc=False):
         print("IDLE Subprocess: no IP port passed in sys.argv.",
               file=sys.__stderr__)
         return
+
+    capture_warnings(True)
     sys.argv[:] = [""]
     sockthread = threading.Thread(target=manage_socket,
                                   name='SockThread',
@@ -118,6 +130,7 @@ def main(del_exitfunc=False):
                 exit_now = True
             continue
         except SystemExit:
+            capture_warnings(False)
             raise
         except:
             type, value, tb = sys.exc_info()
@@ -247,6 +260,7 @@ def exit():
     if no_exitfunc:
         import atexit
         atexit._clear()
+    capture_warnings(False)
     sys.exit(0)
 
 class MyRPCServer(rpc.RPCServer):
@@ -386,3 +400,5 @@ class Executive(object):
         sys.last_value = val
         item = StackViewer.StackTreeItem(flist, tb)
         return RemoteObjectBrowser.remote_object_tree_item(item)
+
+capture_warnings(False)  # Make sure turned off; see issue 18081