Fix #444582. Add shutil.which function for finding programs on the system path.

2 files changed, 93 insertions, 2 deletions

diff --git a/Lib/shutil.py b/Lib/shutil.py
index 46398ef..a4c1436 100644
--- a/Lib/shutil.py
+++ b/Lib/shutil.py
@@ -36,7 +36,7 @@ __all__ = ["copyfileobj", "copyfile", "copymode", "copystat", "copy", "copy2",
            "register_archive_format", "unregister_archive_format",
            "get_unpack_formats", "register_unpack_format",
            "unregister_unpack_format", "unpack_archive",
-           "ignore_patterns", "chown"]
+           "ignore_patterns", "chown", "which"]
            # disk_usage is added later, if available on the platform
 
 class Error(EnvironmentError):
@@ -961,3 +961,51 @@ def get_terminal_size(fallback=(80, 24)):
             lines = size.lines
 
     return os.terminal_size((columns, lines))
+
+def which(cmd, mode=os.F_OK | os.X_OK, path=None):
+    """Given a file, mode, and a path string, return the path whichs conform
+    to the given mode on the path."""
+    # Check that a given file can be accessed with the correct mode.
+    # Additionally check that `file` is not a directory, as on Windows
+    # directories pass the os.access check.
+    def _access_check(fn, mode):
+        if (os.path.exists(fn) and os.access(fn, mode)
+            and not os.path.isdir(fn)):
+            return True
+        return False
+
+    # Short circuit. If we're given a full path which matches the mode
+    # and it exists, we're done here.
+    if _access_check(cmd, mode):
+        return cmd
+
+    path = (path or os.environ.get("PATH", os.defpath)).split(os.pathsep)
+
+    if sys.platform == "win32":
+        # The current directory takes precedence on Windows.
+        if not os.curdir in path:
+            path.insert(0, os.curdir)
+
+        # PATHEXT is necessary to check on Windows.
+        pathext = os.environ.get("PATHEXT", "").split(os.pathsep)
+        # See if the given file matches any of the expected path extensions.
+        # This will allow us to short circuit when given "python.exe".
+        matches = [cmd for ext in pathext if cmd.lower().endswith(ext.lower())]
+        # If it does match, only test that one, otherwise we have to try others.
+        files = [cmd + ext.lower() for ext in pathext] if not matches else [cmd]
+    else:
+        # On other platforms you don't have things like PATHEXT to tell you
+        # what file suffixes are executable, so just pass on cmd as-is.
+        files = [cmd]
+
+    seen = set()
+    for dir in path:
+        dir = os.path.normcase(os.path.abspath(dir))
+        if not dir in seen:
+            seen.add(dir)
+            for thefile in files:
+                name = os.path.join(dir, thefile)
+                if _access_check(name, mode):
+                    return name
+    return None
+
diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
index 1181520..96084ec 100644
--- a/Lib/test/test_shutil.py
+++ b/Lib/test/test_shutil.py
@@ -1128,6 +1128,49 @@ class TestShutil(unittest.TestCase):
         self.assertEqual(['foo'], os.listdir(rv))
 
 
+class TestWhich(unittest.TestCase):
+
+    def setUp(self):
+        self.temp_dir = tempfile.mkdtemp()
+        # Give the temp_file an ".exe" suffix for all.
+        # It's needed on Windows and not harmful on other platforms.
+        self.temp_file = tempfile.NamedTemporaryFile(dir=self.temp_dir,
+                                                     suffix=".exe")
+        os.chmod(self.temp_file.name, stat.S_IXUSR)
+        self.addCleanup(self.temp_file.close)
+        self.dir, self.file = os.path.split(self.temp_file.name)
+
+    def test_basic(self):
+        # Given an EXE in a directory, it should be returned.
+        rv = shutil.which(self.file, path=self.dir)
+        self.assertEqual(rv, self.temp_file.name)
+
+    def test_full_path_short_circuit(self):
+        # When given the fully qualified path to an executable that exists,
+        # it should be returned.
+        rv = shutil.which(self.temp_file.name, path=self.temp_dir)
+        self.assertEqual(self.temp_file.name, rv)
+
+    def test_non_matching_mode(self):
+        # Set the file read-only and ask for writeable files.
+        os.chmod(self.temp_file.name, stat.S_IREAD)
+        rv = shutil.which(self.file, path=self.dir, mode=os.W_OK)
+        self.assertIsNone(rv)
+
+    def test_nonexistent_file(self):
+        # Return None when no matching executable file is found on the path.
+        rv = shutil.which("foo.exe", path=self.dir)
+        self.assertIsNone(rv)
+
+    @unittest.skipUnless(sys.platform == "win32",
+                         "pathext check is Windows-only")
+    def test_pathext_checking(self):
+        # Ask for the file without the ".exe" extension, then ensure that
+        # it gets found properly with the extension.
+        rv = shutil.which(self.temp_file.name[:-4], path=self.dir)
+        self.assertEqual(self.temp_file.name, rv)
+
+
 class TestMove(unittest.TestCase):
 
     def setUp(self):
@@ -1460,7 +1503,7 @@ class TermsizeTests(unittest.TestCase):
 
 def test_main():
     support.run_unittest(TestShutil, TestMove, TestCopyFile,
-                         TermsizeTests)
+                         TermsizeTests, TestWhich)
 
 if __name__ == '__main__':
     test_main()