Add C implementation of ntpath.splitroot()

zooba · Apr 25, 2024 · Apr 19, 2024 · Apr 19, 2024 · Apr 19, 2024 · Apr 19, 2024
commit c397a23922d7d2e7beabee4d834689cfebe1b2b8
diff --git a/Include/internal/pycore_fileutils.h b/Include/internal/pycore_fileutils.h
@@ -290,6 +290,10 @@ extern wchar_t *_Py_normpath_and_size(wchar_t *path, Py_ssize_t size, Py_ssize_t
 extern HRESULT PathCchSkipRoot(const wchar_t *pszPath, const wchar_t **ppszRootEnd);
 #endif /* defined(MS_WINDOWS_GAMES) && !defined(MS_WINDOWS_DESKTOP) */
 
+#ifdef MS_WINDOWS
+extern void _Py_skiproot(wchar_t *path, Py_ssize_t size, Py_ssize_t *drvsize, Py_ssize_t *rootsize);
+#endif
+
 // Macros to protect CRT calls against instant termination when passed an
 // invalid parameter (bpo-23524). IPH stands for Invalid Parameter Handler.
 // Usage:

diff --git a/Lib/ntpath.py b/Lib/ntpath.py
@@ -166,57 +166,76 @@ def splitdrive(p):
     drive, root, tail = splitroot(p)
     return drive, root + tail
 
-
-def splitroot(p):
-    """Split a pathname into drive, root and tail. The drive is defined
-    exactly as in splitdrive(). On Windows, the root may be a single path
-    separator or an empty string. The tail contains anything after the root.
-    For example:
-
-        splitroot('//server/share/') == ('//server/share', '/', '')
-        splitroot('C:/Users/Barney') == ('C:', '/', 'Users/Barney')
-        splitroot('C:///spam///ham') == ('C:', '/', '//spam///ham')
-        splitroot('Windows/notepad') == ('', '', 'Windows/notepad')
-    """
-    p = os.fspath(p)
-    if isinstance(p, bytes):
-        sep = b'\\'
-        altsep = b'/'
-        colon = b':'
-        unc_prefix = b'\\\\?\\UNC\\'
-        empty = b''
-    else:
-        sep = '\\'
-        altsep = '/'
-        colon = ':'
-        unc_prefix = '\\\\?\\UNC\\'
-        empty = ''
-    normp = p.replace(altsep, sep)
-    if normp[:1] == sep:
-        if normp[1:2] == sep:
-            # UNC drives, e.g. \\server\share or \\?\UNC\server\share
-            # Device drives, e.g. \\.\device or \\?\device
-            start = 8 if normp[:8].upper() == unc_prefix else 2
-            index = normp.find(sep, start)
-            if index == -1:
-                return p, empty, empty
-            index2 = normp.find(sep, index + 1)
-            if index2 == -1:
-                return p, empty, empty
-            return p[:index2], p[index2:index2 + 1], p[index2 + 1:]
+try:
+    from nt import _path_splitroot_ex
+except ImportError:
+    def splitroot(p):
+        """Split a pathname into drive, root and tail. The drive is defined
+        exactly as in splitdrive(). On Windows, the root may be a single path
+        separator or an empty string. The tail contains anything after the root.
+        For example:
+
+            splitroot('//server/share/') == ('//server/share', '/', '')
+            splitroot('C:/Users/Barney') == ('C:', '/', 'Users/Barney')
+            splitroot('C:///spam///ham') == ('C:', '/', '//spam///ham')
+            splitroot('Windows/notepad') == ('', '', 'Windows/notepad')
+        """
+        p = os.fspath(p)
+        if isinstance(p, bytes):
+            sep = b'\\'
+            altsep = b'/'
+            colon = b':'
+            unc_prefix = b'\\\\?\\UNC\\'
+            empty = b''
         else:
-            # Relative path with root, e.g. \Windows
-            return empty, p[:1], p[1:]
-    elif normp[1:2] == colon:
-        if normp[2:3] == sep:
-            # Absolute drive-letter path, e.g. X:\Windows
-            return p[:2], p[2:3], p[3:]
+            sep = '\\'
+            altsep = '/'
+            colon = ':'
+            unc_prefix = '\\\\?\\UNC\\'
+            empty = ''
+        normp = p.replace(altsep, sep)
+        if normp[:1] == sep:
+            if normp[1:2] == sep:
+                # UNC drives, e.g. \\server\share or \\?\UNC\server\share
+                # Device drives, e.g. \\.\device or \\?\device
+                start = 8 if normp[:8].upper() == unc_prefix else 2
+                index = normp.find(sep, start)
+                if index == -1:
+                    return p, empty, empty
+                index2 = normp.find(sep, index + 1)
+                if index2 == -1:
+                    return p, empty, empty
+                return p[:index2], p[index2:index2 + 1], p[index2 + 1:]
+            else:
+                # Relative path with root, e.g. \Windows
+                return empty, p[:1], p[1:]
+        elif normp[1:2] == colon:
+            if normp[2:3] == sep:
+                # Absolute drive-letter path, e.g. X:\Windows
+                return p[:2], p[2:3], p[3:]
+            else:
+                # Relative path with drive, e.g. X:Windows
+                return p[:2], empty, p[2:]
         else:
-            # Relative path with drive, e.g. X:Windows
-            return p[:2], empty, p[2:]
-    else:
-        # Relative path, e.g. Windows
-        return empty, empty, p
+            # Relative path, e.g. Windows
+            return empty, empty, p
+else:
+    def splitroot(p):
+        """Split a pathname into drive, root and tail. The drive is defined
+        exactly as in splitdrive(). On Windows, the root may be a single path
+        separator or an empty string. The tail contains anything after the root.
+        For example:
+
+            splitroot('//server/share/') == ('//server/share', '/', '')
+            splitroot('C:/Users/Barney') == ('C:', '/', 'Users/Barney')
+            splitroot('C:///spam///ham') == ('C:', '/', '//spam///ham')
+            splitroot('Windows/notepad') == ('', '', 'Windows/notepad')
+        """
+        p = os.fspath(p)
+        if isinstance(p, bytes):
+            drive, root, tail = _path_splitroot_ex(os.fsdecode(p))
+            return os.fsencode(drive), os.fsencode(root), os.fsencode(tail)
+        return _path_splitroot_ex(p)
 
 
 # Split a path in head (everything up to the last '/') and tail (the

diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -5085,6 +5085,34 @@ os__path_splitroot_impl(PyObject *module, path_t *path)
     return result;
 }
 
+#ifdef MS_WINDOWS
+/*[clinic input]
+os._path_splitroot_ex
+
+    path: unicode
+    /
+
+[clinic start generated code]*/
+
+static PyObject *
+os__path_splitroot_ex_impl(PyObject *module, PyObject *path)
+/*[clinic end generated code: output=de97403d3dfebc40 input=bebce42edb41f967]*/
+{
+    Py_ssize_t len;
+    wchar_t *buffer = PyUnicode_AsWideCharString(path, &len);
+    Py_ssize_t drvsize;
+    Py_ssize_t rootsize;
+    _Py_skiproot(buffer, len, &drvsize, &rootsize);
+    wchar_t *p = buffer;
+    PyObject *drv = PyUnicode_FromWideChar(p, drvsize);
+    p += drvsize;
+    PyObject *root = PyUnicode_FromWideChar(p, rootsize);
+    p += rootsize;
+    PyObject *tail = PyUnicode_FromWideChar(p, len - drvsize - rootsize);
+    PyMem_Free(buffer);
+    return Py_BuildValue("(OOO)", drv, root, tail);
+}
+#endif
 
 /*[clinic input]
 os._path_isdir
@@ -16799,6 +16827,7 @@ static PyMethodDef posix_methods[] = {
     OS__FINDFIRSTFILE_METHODDEF
     OS__GETVOLUMEPATHNAME_METHODDEF
     OS__PATH_SPLITROOT_METHODDEF
+    OS__PATH_SPLITROOT_EX_METHODDEF
     OS__PATH_NORMPATH_METHODDEF
     OS_GETLOADAVG_METHODDEF
     OS_URANDOM_METHODDEF

diff --git a/Python/fileutils.c b/Python/fileutils.c
@@ -2295,6 +2295,63 @@ PathCchCombineEx(wchar_t *buffer, size_t bufsize, const wchar_t *dirname,
 
 #endif /* defined(MS_WINDOWS_GAMES) && !defined(MS_WINDOWS_DESKTOP) */
 
+#ifdef MS_WINDOWS
+void
+_Py_skiproot(wchar_t *path, Py_ssize_t size, Py_ssize_t *drvsize, Py_ssize_t *rootsize)
+{
+    wchar_t *pEnd = size >= 0 ? &path[size] : NULL;
+#define IS_END(x) (pEnd ? (x) == pEnd : !*(x))
+#define IS_SEP(x) (*(x) == SEP || *(x) == ALTSEP)
+    *drvsize = 0;
+    *rootsize = 0;
+    if (IS_SEP(&path[0])) {
+        if (IS_SEP(&path[1])) {
+            // Device drives, e.g. \\.\device or \\?\device
+            // UNC drives, e.g. \\server\share or \\?\UNC\server\share
+            Py_ssize_t idx;
+            if (path[2] == L'?' && IS_SEP(&path[3]) && (path[4] == L'U' || path[4] == L'u') &&
+                (path[5] == L'N' || path[5] == L'n') && (path[6] == L'C' || path[6] == L'c') &&
+                IS_SEP(&path[7])) {
+                idx = 8;
+            }
+            else {
+                idx = 2;
+            }
+            while (!IS_END(&path[idx]) && !IS_SEP(&path[idx])) {
+                idx++;
+            }
+            if (IS_END(&path[idx])) {
+                *drvsize = idx;
+            }
+            else {
+                idx++;
+                while (!IS_END(&path[idx]) && !IS_SEP(&path[idx])) {
+                    idx++;
+                }
+                *drvsize = idx;
+                if (IS_SEP(&path[idx])) {
+                    *rootsize = 1;
+                }
+            }
+        }
+        else {
+            // Relative path with root, e.g. \Windows
+            *rootsize = 1;
+        }
+    }
+    else if (!IS_END(&path[0]) && path[1] == L':')
+    {
+        *drvsize = 2;
+        if (IS_SEP(&path[2])) {
+            // Absolute drive-letter path, e.g. X:\Windows
+            *rootsize = 1;
+        }
+    }
+#undef IS_SEP
+#undef IS_END
+}
+#endif
+
 // The caller must ensure "buffer" is big enough.
 static int
 join_relfile(wchar_t *buffer, size_t bufsize,