bpo-23205: IDLE: Add tests and refactor grep's findfiles (GH-12203)

* Add tests for grep findfiles.
* Move findfiles to module function.
* Change findfiles to use os.walk.

Based on a patch by Al Sweigart.

2 files changed, 107 insertions, 45 deletions

diff --git a/Lib/idlelib/grep.py b/Lib/idlelib/grep.py
index 6068d7e..1251359 100644
--- a/Lib/idlelib/grep.py
+++ b/Lib/idlelib/grep.py
@@ -40,6 +40,27 @@ def grep(text, io=None, flist=None):
     dialog.open(text, searchphrase, io)
 
 
+def walk_error(msg):
+    "Handle os.walk error."
+    print(msg)
+
+
+def findfiles(folder, pattern, recursive):
+    """Generate file names in dir that match pattern.
+
+    Args:
+        folder: Root directory to search.
+        pattern: File pattern to match.
+        recursive: True to include subdirectories.
+    """
+    for dirpath, _, filenames in os.walk(folder, onerror=walk_error):
+        yield from (os.path.join(dirpath, name)
+                    for name in filenames
+                    if fnmatch.fnmatch(name, pattern))
+        if not recursive:
+            break
+
+
 class GrepDialog(SearchDialogBase):
     "Dialog for searching multiple files."
 
@@ -140,15 +161,16 @@ class GrepDialog(SearchDialogBase):
             prog: The compiled, cooked search pattern.
             path: String containing the search path.
         """
-        dir, base = os.path.split(path)
-        list = self.findfiles(dir, base, self.recvar.get())
-        list.sort()
+        folder, filepat = os.path.split(path)
+        if not folder:
+            folder = os.curdir
+        filelist = sorted(findfiles(folder, filepat, self.recvar.get()))
         self.close()
         pat = self.engine.getpat()
         print(f"Searching {pat!r} in {path} ...")
         hits = 0
         try:
-            for fn in list:
+            for fn in filelist:
                 try:
                     with open(fn, errors='replace') as f:
                         for lineno, line in enumerate(f, 1):
@@ -166,36 +188,6 @@ class GrepDialog(SearchDialogBase):
             # so in OW.write, OW.text.insert fails.
             pass
 
-    def findfiles(self, dir, base, rec):
-        """Return list of files in the dir that match the base pattern.
-
-        Use the current directory if dir has no value.
-        If rec is True, recursively iterate through subdirectories.
-
-        Args:
-            dir: Directory path to search.
-            base: File search pattern.
-            rec: Boolean for recursive search through subdirectories.
-        """
-        try:
-            names = os.listdir(dir or os.curdir)
-        except OSError as msg:
-            print(msg)
-            return []
-        list = []
-        subdirs = []
-        for name in names:
-            fn = os.path.join(dir, name)
-            if os.path.isdir(fn):
-                subdirs.append(fn)
-            else:
-                if fnmatch.fnmatch(name, base):
-                    list.append(fn)
-        if rec:
-            for subdir in subdirs:
-                list.extend(self.findfiles(subdir, base, rec))
-        return list
-
 
 def _grep_dialog(parent):  # htest #
     from tkinter import Toplevel, Text, SEL, END
diff --git a/Lib/idlelib/idle_test/test_grep.py b/Lib/idlelib/idle_test/test_grep.py
index ab0d786..a0b5b69 100644
--- a/Lib/idlelib/idle_test/test_grep.py
+++ b/Lib/idlelib/idle_test/test_grep.py
@@ -5,10 +5,11 @@ An exception raised in one method will fail callers.
 Otherwise, tests are mostly independent.
 Currently only test grep_it, coverage 51%.
 """
-from idlelib.grep import GrepDialog
+from idlelib import grep
 import unittest
 from test.support import captured_stdout
 from idlelib.idle_test.mock_tk import Var
+import os
 import re
 
 
@@ -26,23 +27,92 @@ searchengine = Dummy_searchengine()
 class Dummy_grep:
     # Methods tested
     #default_command = GrepDialog.default_command
-    grep_it = GrepDialog.grep_it
-    findfiles = GrepDialog.findfiles
+    grep_it = grep.GrepDialog.grep_it
     # Other stuff needed
     recvar = Var(False)
     engine = searchengine
     def close(self):  # gui method
         pass
 
-grep = Dummy_grep()
+_grep = Dummy_grep()
 
 
 class FindfilesTest(unittest.TestCase):
-    # findfiles is really a function, not a method, could be iterator
-    # test that filename return filename
-    # test that idlelib has many .py files
-    # test that recursive flag adds idle_test .py files
-    pass
+
+    @classmethod
+    def setUpClass(cls):
+        cls.realpath = os.path.realpath(__file__)
+        cls.path = os.path.dirname(cls.realpath)
+
+    @classmethod
+    def tearDownClass(cls):
+        del cls.realpath, cls.path
+
+    def test_invaliddir(self):
+        with captured_stdout() as s:
+            filelist = list(grep.findfiles('invaliddir', '*.*', False))
+        self.assertEqual(filelist, [])
+        self.assertIn('invalid', s.getvalue())
+
+    def test_curdir(self):
+        # Test os.curdir.
+        ff = grep.findfiles
+        save_cwd = os.getcwd()
+        os.chdir(self.path)
+        filename = 'test_grep.py'
+        filelist = list(ff(os.curdir, filename, False))
+        self.assertIn(os.path.join(os.curdir, filename), filelist)
+        os.chdir(save_cwd)
+
+    def test_base(self):
+        ff = grep.findfiles
+        readme = os.path.join(self.path, 'README.txt')
+
+        # Check for Python files in path where this file lives.
+        filelist = list(ff(self.path, '*.py', False))
+        # This directory has many Python files.
+        self.assertGreater(len(filelist), 10)
+        self.assertIn(self.realpath, filelist)
+        self.assertNotIn(readme, filelist)
+
+        # Look for .txt files in path where this file lives.
+        filelist = list(ff(self.path, '*.txt', False))
+        self.assertNotEqual(len(filelist), 0)
+        self.assertNotIn(self.realpath, filelist)
+        self.assertIn(readme, filelist)
+
+        # Look for non-matching pattern.
+        filelist = list(ff(self.path, 'grep.*', False))
+        self.assertEqual(len(filelist), 0)
+        self.assertNotIn(self.realpath, filelist)
+
+    def test_recurse(self):
+        ff = grep.findfiles
+        parent = os.path.dirname(self.path)
+        grepfile = os.path.join(parent, 'grep.py')
+        pat = '*.py'
+
+        # Get Python files only in parent directory.
+        filelist = list(ff(parent, pat, False))
+        parent_size = len(filelist)
+        # Lots of Python files in idlelib.
+        self.assertGreater(parent_size, 20)
+        self.assertIn(grepfile, filelist)
+        # Without subdirectories, this file isn't returned.
+        self.assertNotIn(self.realpath, filelist)
+
+        # Include subdirectories.
+        filelist = list(ff(parent, pat, True))
+        # More files found now.
+        self.assertGreater(len(filelist), parent_size)
+        self.assertIn(grepfile, filelist)
+        # This file exists in list now.
+        self.assertIn(self.realpath, filelist)
+
+        # Check another level up the tree.
+        parent = os.path.dirname(parent)
+        filelist = list(ff(parent, '*.py', True))
+        self.assertIn(self.realpath, filelist)
 
 
 class Grep_itTest(unittest.TestCase):
@@ -51,9 +121,9 @@ class Grep_itTest(unittest.TestCase):
     # from incomplete replacement, so 'later'.
 
     def report(self, pat):
-        grep.engine._pat = pat
+        _grep.engine._pat = pat
         with captured_stdout() as s:
-            grep.grep_it(re.compile(pat), __file__)
+            _grep.grep_it(re.compile(pat), __file__)
         lines = s.getvalue().split('\n')
         lines.pop()  # remove bogus '' after last \n
         return lines