bpo-41138: Fix trace CLI for non-UTF-8 files. (GH-21177)

Fix also a resource warning when store counts and module info.

2 files changed, 25 insertions, 16 deletions

diff --git a/Lib/test/test_trace.py b/Lib/test/test_trace.py
index 89d46376bff..c03982ba72b 100644
--- a/Lib/test/test_trace.py
+++ b/Lib/test/test_trace.py
@@ -1,6 +1,6 @@
 import os
 import sys
-from test.support import TESTFN, rmtree, unlink, captured_stdout
+from test.support import TESTFN, TESTFN_UNICODE, FS_NONASCII, rmtree, unlink, captured_stdout
 from test.support.script_helper import assert_python_ok, assert_python_failure
 import textwrap
 import unittest
@@ -428,9 +428,10 @@ class TestCoverageCommandLineOutput(unittest.TestCase):
     coverfile = 'tmp.cover'
 
     def setUp(self):
-        with open(self.codefile, 'w') as f:
+        with open(self.codefile, 'w', encoding='iso-8859-15') as f:
             f.write(textwrap.dedent('''\
-                x = 42
+                # coding: iso-8859-15
+                x = 'spœm'
                 if []:
                     print('unreachable')
             '''))
@@ -451,9 +452,10 @@ class TestCoverageCommandLineOutput(unittest.TestCase):
         self.assertEqual(stderr, b'')
         self.assertFalse(os.path.exists(tracecoverpath))
         self.assertTrue(os.path.exists(self.coverfile))
-        with open(self.coverfile) as f:
+        with open(self.coverfile, encoding='iso-8859-15') as f:
             self.assertEqual(f.read(),
-                "    1: x = 42\n"
+                "       # coding: iso-8859-15\n"
+                "    1: x = 'spœm'\n"
                 "    1: if []:\n"
                 "           print('unreachable')\n"
             )
@@ -462,9 +464,10 @@ class TestCoverageCommandLineOutput(unittest.TestCase):
         argv = '-m trace --count --missing'.split() + [self.codefile]
         status, stdout, stderr = assert_python_ok(*argv)
         self.assertTrue(os.path.exists(self.coverfile))
-        with open(self.coverfile) as f:
+        with open(self.coverfile, encoding='iso-8859-15') as f:
             self.assertEqual(f.read(), textwrap.dedent('''\
-                    1: x = 42
+                       # coding: iso-8859-15
+                    1: x = 'spœm'
                     1: if []:
                 >>>>>>     print('unreachable')
             '''))
@@ -485,15 +488,19 @@ class TestCommandLine(unittest.TestCase):
             self.assertIn(message, stderr)
 
     def test_listfuncs_flag_success(self):
-        with open(TESTFN, 'w') as fd:
-            self.addCleanup(unlink, TESTFN)
+        filename = TESTFN + '.py'
+        modulename = os.path.basename(TESTFN)
+        with open(filename, 'w', encoding='utf-8') as fd:
+            self.addCleanup(unlink, filename)
             fd.write("a = 1\n")
-            status, stdout, stderr = assert_python_ok('-m', 'trace', '-l', TESTFN,
+            status, stdout, stderr = assert_python_ok('-m', 'trace', '-l', filename,
                                                       PYTHONIOENCODING='utf-8')
             self.assertIn(b'functions called:', stdout)
+            expected = f'filename: {filename}, modulename: {modulename}, funcname: <module>'
+            self.assertIn(expected.encode(), stdout)
 
     def test_sys_argv_list(self):
-        with open(TESTFN, 'w') as fd:
+        with open(TESTFN, 'w', encoding='utf-8') as fd:
             self.addCleanup(unlink, TESTFN)
             fd.write("import sys\n")
             fd.write("print(type(sys.argv))\n")
@@ -506,7 +513,8 @@ class TestCommandLine(unittest.TestCase):
     def test_count_and_summary(self):
         filename = f'{TESTFN}.py'
         coverfilename = f'{TESTFN}.cover'
-        with open(filename, 'w') as fd:
+        modulename = os.path.basename(TESTFN)
+        with open(filename, 'w', encoding='utf-8') as fd:
             self.addCleanup(unlink, filename)
             self.addCleanup(unlink, coverfilename)
             fd.write(textwrap.dedent("""\
@@ -524,7 +532,7 @@ class TestCommandLine(unittest.TestCase):
         stdout = stdout.decode()
         self.assertEqual(status, 0)
         self.assertIn('lines   cov%   module   (path)', stdout)
-        self.assertIn(f'6   100%   {TESTFN}   ({filename})', stdout)
+        self.assertIn(f'6   100%   {modulename}   ({filename})', stdout)
 
     def test_run_as_module(self):
         assert_python_ok('-m', 'trace', '-l', '--module', 'timeit', '-n', '1')
diff --git a/Lib/trace.py b/Lib/trace.py
index 52047c3fbf4..c505d8bc72a 100755
--- a/Lib/trace.py
+++ b/Lib/trace.py
@@ -287,8 +287,9 @@ class CoverageResults:
         if self.outfile:
             # try and store counts and module info into self.outfile
             try:
-                pickle.dump((self.counts, self.calledfuncs, self.callers),
-                            open(self.outfile, 'wb'), 1)
+                with open(self.outfile, 'wb') as f:
+                    pickle.dump((self.counts, self.calledfuncs, self.callers),
+                                f, 1)
             except OSError as err:
                 print("Can't save counts files because %s" % err, file=sys.stderr)
 
@@ -715,7 +716,7 @@ def main():
             sys.argv = [opts.progname, *opts.arguments]
             sys.path[0] = os.path.dirname(opts.progname)
 
-            with open(opts.progname) as fp:
+            with open(opts.progname, 'rb') as fp:
                 code = compile(fp.read(), opts.progname, 'exec')
             # try to emulate __main__ namespace as much as possible
             globs = {