From 569355c7670c73170293cc671e81f89093b8b505 Mon Sep 17 00:00:00 2001
From: Thor Whalen <thorwhalen1@gmail.com>
Date: Wed, 25 Oct 2023 13:25:20 +0200
Subject: [PATCH] feat: paths.extract_path_segments

---
 xdol/paths.py | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 xdol/paths.py

diff --git a/xdol/paths.py b/xdol/paths.py
new file mode 100644
index 0000000..6db5d49
--- /dev/null
+++ b/xdol/paths.py
@@ -0,0 +1,82 @@
+"""Tools for manipulating paths.
+
+A path is a sequence of objects that represent a reference or location. 
+Often paths are strings that use a special character (e.g. '/') to distinguish the 
+separate objects in the sequence.
+
+"""
+
+import os
+
+
+def extract_path_segments(path, *, sep=os.path.sep, depth=1, trailing_sep=True):
+    """Extracts the first `depth` segments of a path, delimited by `sep`.
+
+    :param path: The path to extract the segments from.
+    :param sep: The delimiter to use to separate the segments.
+    :param depth: The number of segments to extract.
+    :param trailing_sep: Whether to add a trailing delimiter to the result.
+    :return: The first `depth` segments of the path, delimited by `sep`.
+
+    By default, the function will use depth=1 and the system's path separator as `sep`.
+
+    >>> extract_path_segments("snap/crackle/pop/cereal", sep='/')
+    'snap/'
+
+    You can specify a different `depth`.
+
+    >>> extract_path_segments("snap/crackle/pop/cereal", depth=2, sep='/')
+    'snap/crackle/'
+
+    Note that if the path starts with the `sep` delimiter, the result will be different.
+
+    >>> extract_path_segments("/snap/crackle/pop/cereal", depth=2, sep='/')
+    '/snap/'
+
+    If `path` is not a string, the function will assume it's an iterable of strings,
+    and return an iterable of the results.
+
+    >>> list(extract_path_segments(
+    ...     ['a.file', 'a/folder', 'a/folder_with/subfolders/', 'another/folder'],
+    ...     depth=2, sep='/'
+    ... ))
+    ['a.file', 'a/folder', 'a/folder_with/', 'another/folder']
+
+    You can of course specify a different delimiter.
+
+    >>> extract_path_segments("snap.crackle.pop.cereal", sep='.', depth=2)
+    'snap.crackle.'
+
+    That last dot doesn't make sense here? Tell the function to not add the separator
+    at the end of the result.
+
+    >>> extract_path_segments(
+    ...     "snap.crackle.pop.cereal", sep='.', depth=2, trailing_sep=False
+    ... )
+    'snap.crackle'
+
+    """
+    if not isinstance(path, str):
+        from functools import partial
+
+        paths = path  # path is actually an iterable of paths
+        return map(partial(extract_path_segments, sep=sep, depth=depth), paths)
+    else:
+        if depth < 1:
+            raise ValueError("depth must be a positive integer")
+
+        segments = path.split(sep)
+
+        if depth >= len(segments):
+            # If the depth specified exceeds the actual number of segments in the path,
+            # return the entire path as is
+            return path
+
+        # Join the segments up to the specified depth with the delimiter
+        result = sep.join(segments[:depth])
+
+        # Append a trailing delimiter to indicate the last segment
+        if trailing_sep:
+            result += sep
+
+        return result