Add a number of utility functions.

These will be important in future commits.

2 files changed, 151 insertions, 24 deletions

diff --git a/mu4web/main.py b/mu4web/main.py
index 3cc39dc..8a4baec 100644
--- a/mu4web/main.py
+++ b/mu4web/main.py
@@ -46,6 +46,7 @@ from .components import (
     attachement_tree,
     login_page,
 )
+from .util import MutableString
 
 #
 # A few operations depend on the index of attachements. These index
@@ -195,29 +196,6 @@ def raw_message():
     return flask.send_file(filename, mimetype='message/rfc822')
 
 
-class MutableString:
-    """
-    A mutatable string.
-
-    Strings are immutable by default in python. This works almost
-    exactly like a regular string, but ``+=`` actually changes the
-    object in place.
-    """
-
-    def __init__(self) -> None:
-        self.str = ''
-
-    def __iadd__(self, other: str) -> 'MutableString':
-        self.str += other
-        return self
-
-    def __repr__(self) -> str:
-        return f'MutableString("{self.str}")'
-
-    def __str__(self) -> str:
-        return self.str
-
-
 class IMGParser(HTMLParser):
     """
     Rewrites HTML image tags to be safer/have more functionality.
diff --git a/mu4web/util.py b/mu4web/util.py
index 7e2df8c..9df042d 100644
--- a/mu4web/util.py
+++ b/mu4web/util.py
@@ -2,12 +2,25 @@
 
 import subprocess
 from os import PathLike
+from contextlib import contextmanager
+import os
+from dataclasses import dataclass
 from typing import (
+    Callable,
+    Generic,
+    Iterator,
+    TypeVar,
     Union,
+    overload,
+    Literal,
 )
 
+# Ts = TypeVarTuple('Ts')
+T = TypeVar('T')
+V = TypeVar('V')
 
-def find(basedir: PathLike[str] | PathLike[bytes],
+
+def find(basedir: str | bytes | PathLike[str] | PathLike[bytes],
          **flags: str | bytes) -> list[bytes]:
     """
     Run the shell command ``find``.
@@ -29,3 +42,139 @@ def find(basedir: PathLike[str] | PathLike[bytes],
 
     cmd = subprocess.run(cmdline, capture_output=True)
     return cmd.stdout.split(b'\0')[:-1]
+
+
+@contextmanager
+def cwd(path: str) -> Iterator[None]:
+    """
+    Context manager for changing directories.
+
+    Changes the current working directory for the block. And changes
+    it back after, no matter how the block is left.
+
+    .. code-block:: python
+
+        with cwd("/tmp"):
+            print(os.getcwd())
+            # /tmp
+    """
+    oldpwd = os.getcwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(oldpwd)
+
+
+class chain(Generic[T]):
+    """
+    Chain functions after one another.
+
+    .. code-block:: python
+
+        (chain(range(9))
+            @ sum
+            @ sqrt).value
+        # result: 6.0
+
+    Is equivalent to
+
+    .. code-block:: python
+
+        sqrt(sum(range(9)))
+        # result: 6.0
+    """
+
+    def __init__(self, value: T) -> None:
+        self.value = value
+
+    def __matmul__(self, proc: Callable[[T], V]) -> 'chain[V]':
+        return chain(proc(self.value))
+
+
+def force(x: T | None) -> T:
+    """
+    Discard the None case from an optional value.
+
+    Only use when you *know* that the value is there.
+
+    :raises AssertionError:
+        When value is None after all.
+    """
+    assert x is not None
+    return x
+
+
+class MutableString:
+    """
+    A mutatable string.
+
+    Strings are immutable by default in python. This works almost
+    exactly like a regular string, but ``+=`` actually changes the
+    object in place.
+    """
+
+    def __init__(self, init: str = '') -> None:
+        self.str = init
+
+    def __add__(self, other: str) -> 'MutableString':
+        return MutableString(self.str + other)
+
+    def __iadd__(self, other: str) -> 'MutableString':
+        self.str += other
+        return self
+
+    def __repr__(self) -> str:
+        return f'MutableString("{self.str}")'
+
+    def __str__(self) -> str:
+        return self.str
+
+
+@dataclass
+class Lists(Generic[T, V]):
+    """
+    A group of lists.
+
+    In an ideal world, this would take any number of type parameters,
+    and work correctly. Currently it's limited to exactly two lists.
+
+    :param lists:
+        The contained lists
+    """
+
+    def __init__(self, ts: list[T] = [], vs: list[V] = []) -> None:
+        self.lists = (ts, vs)
+
+    def __add__(self, other: 'Lists[T, V]') -> 'Lists[T, V]':
+        """Add each sublist of the two groups."""
+        return Lists(self[0] + other[0],
+                     self[1] + other[1])
+
+    @overload
+    def __getitem__(self, idx: Literal[0]) -> list[T]:
+        ...  # pragma: no cover
+
+    @overload
+    def __getitem__(self, idx: Literal[1]) -> list[V]:
+        ...  # pragma: no cover
+
+    def __getitem__(self, idx: int) -> list[T] | list[V]:
+        return self.lists[idx]
+
+
+P = TypeVar('P', str, bytes)
+
+
+# TODO test this on a system with drive letters
+def split_path(path: P) -> list[P]:
+    """Split a path into all its components."""
+    result: list[P] = []
+    dir, item = os.path.split(path)
+    while True:
+        if not item:
+            break
+        result.insert(0, item)
+        path = dir
+        dir, item = os.path.split(path)
+    return result