gh-101000: Add os.path.splitroot() (#101002)

Co-authored-by: Eryk Sun <eryksun@gmail.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

1 files changed, 32 insertions, 11 deletions

diff --git a/Lib/posixpath.py b/Lib/posixpath.py
index 737f8a5c156..32b5d6e105d 100644
--- a/Lib/posixpath.py
+++ b/Lib/posixpath.py
@@ -28,7 +28,7 @@ import stat
 import genericpath
 from genericpath import *
 
-__all__ = ["normcase","isabs","join","splitdrive","split","splitext",
+__all__ = ["normcase","isabs","join","splitdrive","splitroot","split","splitext",
            "basename","dirname","commonprefix","getsize","getmtime",
            "getatime","getctime","islink","exists","lexists","isdir","isfile",
            "ismount", "expanduser","expandvars","normpath","abspath",
@@ -135,6 +135,35 @@ def splitdrive(p):
     return p[:0], p
 
 
+def splitroot(p):
+    """Split a pathname into drive, root and tail. On Posix, drive is always
+    empty; the root may be empty, a single slash, or two slashes. The tail
+    contains anything after the root. For example:
+
+        splitroot('foo/bar') == ('', '', 'foo/bar')
+        splitroot('/foo/bar') == ('', '/', 'foo/bar')
+        splitroot('//foo/bar') == ('', '//', 'foo/bar')
+        splitroot('///foo/bar') == ('', '/', '//foo/bar')
+    """
+    p = os.fspath(p)
+    if isinstance(p, bytes):
+        sep = b'/'
+        empty = b''
+    else:
+        sep = '/'
+        empty = ''
+    if p[:1] != sep:
+        # Relative path, e.g.: 'foo'
+        return empty, empty, p
+    elif p[1:2] != sep or p[2:3] == sep:
+        # Absolute path, e.g.: '/foo', '///foo', '////foo', etc.
+        return empty, sep, p[1:]
+    else:
+        # Precisely two leading slashes, e.g.: '//foo'. Implementation defined per POSIX, see
+        # https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap04.html#tag_04_13
+        return empty, p[:2], p[2:]
+
+
 # Return the tail (basename) part of a path, same as split(path)[1].
 
 def basename(p):
@@ -372,13 +401,7 @@ except ImportError:
             dotdot = '..'
         if path == empty:
             return dot
-        initial_slashes = path.startswith(sep)
-        # POSIX allows one or two initial slashes, but treats three or more
-        # as single slash.
-        # (see https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap04.html#tag_04_13)
-        if (initial_slashes and
-            path.startswith(sep*2) and not path.startswith(sep*3)):
-            initial_slashes = 2
+        _, initial_slashes, path = splitroot(path)
         comps = path.split(sep)
         new_comps = []
         for comp in comps:
@@ -390,9 +413,7 @@ except ImportError:
             elif new_comps:
                 new_comps.pop()
         comps = new_comps
-        path = sep.join(comps)
-        if initial_slashes:
-            path = sep*initial_slashes + path
+        path = initial_slashes + sep.join(comps)
         return path or dot
 
 else: