misc.python.materialize.ui

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10"""Utilities for interacting with humans."""
 11
 12import asyncio
 13import datetime
 14import os
 15import shlex
 16import sys
 17import time
 18from collections.abc import AsyncGenerator, Callable, Generator, Iterable
 19from contextlib import contextmanager
 20from typing import Any
 21
 22from colored import attr, fg
 23
 24
 25class Verbosity:
 26    """How noisy logs should be"""
 27
 28    quiet: bool = False
 29
 30    @classmethod
 31    def init_from_env(cls, explicit: bool | None) -> None:
 32        """Set to quiet based on MZ_QUIET being set to almost any value
 33
 34        The only values that this gets set to false for are the empty string, 0, or no
 35        """
 36        cls.quiet = env_is_truthy("MZ_QUIET")
 37        if explicit is not None:
 38            cls.quiet = explicit
 39
 40
 41def speaker(prefix: str) -> Callable[..., None]:
 42    """Create a function that will log with a prefix to stderr.
 43
 44    Obeys `Verbosity.quiet`. Note that you must include any necessary
 45    spacing after the prefix.
 46
 47    Example::
 48
 49        >>> say = speaker("mz> ")
 50        >>> say("hello")  # doctest: +SKIP
 51        mz> hello
 52    """
 53
 54    def say(msg: str) -> None:
 55        if not Verbosity.quiet:
 56            print(f"{prefix}{msg}", file=sys.stderr)
 57
 58    return say
 59
 60
 61header = speaker("==> ")
 62section = speaker("--- ")
 63say = speaker("")
 64
 65
 66def warn(message: str) -> None:
 67    """Emits a warning message to stderr."""
 68    print(f"{fg('yellow')}warning:{attr('reset')} {message}")
 69
 70
 71def confirm(question: str) -> bool:
 72    """Render a question, returning True if the user says y or yes"""
 73    response = input(f"{question} [y/N]")
 74    return response.lower() in ("y", "yes")
 75
 76
 77def progress(msg: str = "", prefix: str | None = None, *, finish: bool = False) -> None:
 78    """Print a progress message to stderr, using the same prefix format as speaker"""
 79    if prefix is not None:
 80        msg = f"{prefix}> {msg}"
 81    end = "" if not finish else "\n"
 82    print(msg, file=sys.stderr, flush=True, end=end)
 83
 84
 85def timeout_loop(timeout: int, tick: float = 1.0) -> Generator[float, None, None]:
 86    """Loop until timeout, optionally sleeping until tick
 87
 88    Always iterates at least once
 89
 90    Args:
 91        timeout: maximum. number of seconds to wait
 92        tick: how long to ensure passes between loop iterations. Default: 1
 93    """
 94    end = time.monotonic() + timeout
 95    while True:
 96        before = time.monotonic()
 97        yield end - before
 98        after = time.monotonic()
 99
100        if after >= end:
101            return
102
103        if after - before < tick:
104            if tick > 0:
105                time.sleep(tick - (after - before))
106
107
108async def async_timeout_loop(
109    timeout: int, tick: float = 1.0
110) -> AsyncGenerator[float, None]:
111    """Loop until timeout, asynchronously sleeping until tick
112
113    Always iterates at least once
114
115    Args:
116        timeout: maximum. number of seconds to wait
117        tick: how long to ensure passes between loop iterations. Default: 1
118    """
119    end = time.monotonic() + timeout
120    while True:
121        before = time.monotonic()
122        yield end - before
123        after = time.monotonic()
124
125        if after >= end:
126            return
127
128        if after - before < tick:
129            if tick > 0:
130                await asyncio.sleep(tick - (after - before))
131
132
133def log_in_automation(msg: str) -> None:
134    """Log to a file, if we're running in automation"""
135    if env_is_truthy("MZ_IN_AUTOMATION"):
136        with open("/tmp/mzcompose.log", "a") as fh:
137            now = datetime.datetime.now().isoformat()
138            print(f"[{now}] {msg}", file=fh)
139
140
141def shell_quote(args: Iterable[Any]) -> str:
142    """Return shell-escaped string of all the parameters
143
144    ::
145
146        >>> shell_quote(["one", "two three"])
147        "one 'two three'"
148    """
149    return " ".join(shlex.quote(str(arg)) for arg in args)
150
151
152def env_is_truthy(env_var: str, default: str = "0") -> bool:
153    """Return true if `env_var` is set and is not one of: 0, '', no, false"""
154    env = os.getenv(env_var, default)
155    if env is not None:
156        return env not in ("", "0", "no", "false")
157    return False
158
159
160class UIError(Exception):
161    """An error intended for display to humans.
162
163    Use this exception type when the error is something the user can be expected
164    to handle. If the error indicates a truly unexpected condition (i.e., a
165    programming error), use a different exception type that will produce a
166    backtrace instead.
167
168    Attributes:
169        hint: An optional hint to display alongside the error message.
170    """
171
172    def __init__(self, message: str, hint: str | None = None):
173        super().__init__(message)
174        self.hint = hint
175
176    def set_hint(self, hint: str) -> None:
177        """Attaches a hint to the error.
178
179        This method will overwrite the existing hint, if any.
180        """
181        self.hint = hint
182
183
184class CommandFailureCausedUIError(UIError):
185    """
186    An UIError that is caused by executing a command.
187    """
188
189    def __init__(
190        self,
191        message: str,
192        cmd: list[str],
193        stdout: str | None = None,
194        stderr: str | None = None,
195        hint: str | None = None,
196    ):
197        super().__init__(message, hint)
198        self.cmd = cmd
199        self.stdout = stdout
200        self.stderr = stderr
201
202
203@contextmanager
204def error_handler(prog: str) -> Any:
205    """Catches and pretty-prints any raised `UIError`s.
206
207    Args:
208        prog: The name of the program with which to prefix the error message.
209    """
210    try:
211        yield
212    except UIError as e:
213        print(f"{prog}: {fg('red')}error:{attr('reset')} {e}", file=sys.stderr)
214        if e.hint:
215            print(f"{attr('bold')}hint:{attr('reset')} {e.hint}")
216        sys.exit(1)
217    except KeyboardInterrupt:
218        sys.exit(1)