aboutsummaryrefslogtreecommitdiff
path: root/mu4web/util.py
blob: 9df042d88b41f387cb1d512a0f1f15ab553f6395 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""Various misc. utilities."""

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: str | bytes | PathLike[str] | PathLike[bytes],
         **flags: str | bytes) -> list[bytes]:
    """
    Run the shell command ``find``.

    :param basedir:
        Directory to search under.

    :param flags:
        Key-value pairs passed to find. Each key is prefixed with a
        single dash. Compound searches might be possible with some
        trickery, but extend this function instead of doing that.
    """
    cmdline: list[Union[str, bytes,
                        PathLike[str],
                        PathLike[bytes]]] = ['find', basedir]
    for key, value in flags.items():
        cmdline += [f'-{key}', value]
    cmdline.append('-print0')

    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