[stdlib] Add os.path.splitroot

docs/changelog.md

@@ -232,8 +232,10 @@ what we publish.
   of variables that are handled as synthetic types, e.g. `List` from Mojo or
   `std::vector` from C++.
 
-- Added `os.path.expandvars` to expand environment variables in a string.
-  ([PR #3735](https://github.com/modularml/mojo/pull/3735) by [@thatstoasty](https://github.com/thatstoasty)).
+- Expanded `os.path` with new functions (by [@thatstoasty](https://github.com/thatstoasty)):
+  - `os.path.expandvars`: Expands environment variables in a path ([PR #3735](https://github.com/modularml/mojo/pull/3735)).
+  - `os.path.splitroot`: Split a path into drive, root and tail.
+  ([PR #3780](https://github.com/modularml/mojo/pull/3780)).
 
 - Added a `reserve` method and new constructor to the `String` struct to
   allocate additional capacity.

stdlib/src/os/path/__init__.mojo

@@ -22,5 +22,6 @@ from .path import (
     islink,
     join,
     split,
+    splitroot,
     lexists,
 )

stdlib/src/os/path/path.mojo

@@ -392,6 +392,41 @@ def split[PathLike: os.PathLike, //](path: PathLike) -> (String, String):
 
 #     return join(path.__fspath__(), *paths_str)
 
+# ===----------------------------------------------------------------------=== #
+# splitroot
+# ===----------------------------------------------------------------------=== #
+
+
+fn splitroot[
+    PathLike: os.PathLike, //
+](path: PathLike) -> Tuple[String, String, String]:
+    """Splits `path` into drive, root and tail. The tail contains anything after the root.
+
+    Parameters:
+        PathLike: The type conforming to the os.PathLike trait.
+
+    Args:
+        path: The path to be split.
+
+    Returns:
+        A tuple containing three strings: (drive, root, tail).
+    """
+    var p = path.__fspath__()
+    alias empty = String("")
+
+    # Relative path, e.g.: 'foo'
+    if p[:1] != sep:
+        return empty, empty, p
+
+    # Absolute path, e.g.: '/foo', '///foo', '////foo', etc.
+    elif p[1:2] != sep or p[2:3] == sep:
+        return empty, String(sep), p[1:]
+
+    # 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
+    else:
+        return empty, p[:2], p[2:]
+
 
 # ===----------------------------------------------------------------------=== #
 # expandvars

stdlib/test/os/path/test_splitroot.mojo

@@ -0,0 +1,87 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+# RUN: %mojo %s
+
+import os
+from os.path import splitroot
+from testing import assert_equal
+
+
+def test_absolute_path():
+    drive, root, tail = splitroot("/usr/lib/file.txt")
+    assert_equal(drive, "")
+    assert_equal(root, "/")
+    assert_equal(tail, "usr/lib/file.txt")
+
+    drive, root, tail = splitroot("//usr/lib/file.txt")
+    assert_equal(drive, "")
+    assert_equal(root, "//")
+    assert_equal(tail, "usr/lib/file.txt")
+
+    drive, root, tail = splitroot("///usr/lib/file.txt")
+    assert_equal(drive, "")
+    assert_equal(root, "/")
+    assert_equal(tail, "//usr/lib/file.txt")
+
+
+def test_relative_path():
+    drive, root, tail = splitroot("usr/lib/file.txt")
+    assert_equal(drive, "")
+    assert_equal(root, "")
+    assert_equal(tail, "usr/lib/file.txt")
+
+    drive, root, tail = splitroot(".")
+    assert_equal(drive, "")
+    assert_equal(root, "")
+    assert_equal(tail, ".")
+
+    drive, root, tail = splitroot("..")
+    assert_equal(drive, "")
+    assert_equal(root, "")
+    assert_equal(tail, "..")
+
+    drive, root, tail = splitroot("entire/.//.tail/..//captured////")
+    assert_equal(drive, "")
+    assert_equal(root, "")
+    assert_equal(tail, "entire/.//.tail/..//captured////")
+
+
+def test_root_directory():
+    drive, root, tail = splitroot("/")
+    assert_equal(drive, "")
+    assert_equal(root, "/")
+    assert_equal(tail, "")
+
+    drive, root, tail = splitroot("//")
+    assert_equal(drive, "")
+    assert_equal(root, "//")
+    assert_equal(tail, "")
+
+    drive, root, tail = splitroot("///")
+    assert_equal(drive, "")
+    assert_equal(root, "/")
+    assert_equal(tail, "//")
+
+
+def test_empty_path():
+    drive, root, tail = splitroot("")
+    assert_equal(drive, "")
+    assert_equal(root, "")
+    assert_equal(tail, "")
+
+
+def main():
+    test_absolute_path()
+    test_relative_path()
+    test_root_directory()
+    test_empty_path()