my .pythonrc (2024.06.23)

.pythonrc.py

#!/usr/bin/env python3
# pyright: reportUnusedCallResult = false, reportUnusedImport = false
# ruff: noqa: ANN101, D105, D107, D202, D205, D212, D401, DTZ005, EM101, S311, T201, TRY003

"""
.pythonrc is a file which is executed when the Python interactive shell is
started if $PYTHONSTARTUP is in your environment and points to this file. It's
just regular Python code, so do what you will. Your ~/.inputrc file can
greatly complement this file.

Modified from sontek's dotfiles repo on github:
https://github.com/sontek/dotfiles
"""

## LINT COMMAND ###################################################
# ruff check ~/.pythonrc --select ALL --ignore D212 --ignore D203 #
###################################################################

from __future__ import annotations

import abc
import ast
import atexit
import builtins
import calendar as _calendar
import collections.abc
import copy
import ctypes  # noqa: F401
import dataclasses
import datetime
import faulthandler
import functools  # noqa: F401
import importlib.util
import inspect
import io
import math  # noqa: F401
import operator  # noqa: F401
import os
import platform
import random
import re
import subprocess
import sys
import textwrap
import traceback
import typing

import rich.console
from rich import print

if typing.TYPE_CHECKING:
    import types

_PRIDE_MONTH_NUMBER = 6

__MIN_RECURSION_LIMIT = 0x800
__MAX_RECURSION_LIMIT = 0x8000

#############
# LAUNCHING #
#############

LAUNCHING_MESSAGE = (
    "\x1b[35m◉ \x1b[1mLaunching the custom REPL...\x1b[22;39m\n[session]\n"
)
VIRTUAL_ENV = os.environ.get("VIRTUAL_ENV", None)
HOME = VIRTUAL_ENV or os.environ.get("WORKON_HOME", None) or os.environ["HOME"]


#############
# TYPE VARS #
#############

AnyCallable: typing.TypeAlias = collections.abc.Callable[..., typing.Any]
AnyCallableT = typing.TypeVar("AnyCallableT", bound=AnyCallable)

_T0 = typing.TypeVar("_T0")
_T1 = typing.TypeVar("_T1")
_T2 = typing.TypeVar("_T2")


#################################
# SAVE & RESTORE HISTORY STATES #
#################################

try:
    import readline
except ImportError:
    pass
else:
    ##################
    # TAB COMPLETION #
    ##################

    try:
        import rlcompleter  # noqa: F401
    except ImportError:
        pass
    else:
        if sys.platform == "darwin":
            # Work around a bug in Mac OS X's readline module.
            readline.parse_and_bind("bind ^I rl_complete")
        else:
            readline.parse_and_bind("tab: complete")

    ######################
    # PERSISTENT HISTORY #
    ######################

    # Use separate history files for each virtual environment.
    HISTORY_FILE = os.path.join(HOME, ".pyhistory")  # noqa: PTH118

    # Read the existing history if there is one.
    if os.path.exists(HISTORY_FILE):  # noqa: PTH110
        try:
            readline.read_history_file(HISTORY_FILE)
        except Exception:  # noqa: BLE001
            # If there was a problem reading the history file then it may have
            # become corrupted, so we just delete it.
            os.remove(HISTORY_FILE)  # noqa: PTH107

    # Set maximum number of commands written to the history file.
    readline.set_history_length(256)

    @atexit.register
    def savehist() -> None:
        """
        Save the history of the shell into a `.pyhistory` file.

        Automatically runs when the user exists the shell.
        """
        try:
            readline.write_history_file(HISTORY_FILE)
        except NameError:
            pass
        except Exception as err:  # noqa: BLE001
            print(
                f"Unable to save history file due to the following error: {err}",
                file=sys.stderr,
            )


#################
# COLOR SUPPORT #
#################


class TermColors(dict[str, str]):
    """
    Gives easy access to ANSI color codes.

    Attempts to fall back to no color for certain TERM values.
    Mostly taken from IPython.
    """

    COLOR_TEMPLATES = (
        ("Black", "30"),
        ("Red", "31"),
        ("Green", "32"),
        ("Brown", "33"),
        ("Blue", "34"),
        ("Purple", "35"),
        ("Cyan", "36"),
        ("LightGray", "37"),
        ("DarkGray", "90"),
        ("LightRed", "91"),
        ("LightGreen", "92"),
        ("Yellow", "93"),
        ("LightBlue", "94"),
        ("LightPurple", "95"),
        ("LightCyan", "96"),
        ("White", "97"),
        ("Normal", "39"),
    )

    NoColor = ""
    _base = "\033[%sm"

    def __init__(self) -> None:
        if os.environ.get("TERM") in (
            "xterm-color",
            "xterm-256color",
            "linux",
            "screen",
            "screen-256color",
            "screen-bce",
        ):
            self.update((k, self._base % v) for k, v in self.COLOR_TEMPLATES)
        else:
            self.update((k, self.NoColor) for k, _ in self.COLOR_TEMPLATES)


_c: typing.Final = TermColors()


#################
# FAKE COMMANDS #
#################


class FakeCommand(abc.ABC):
    """Class to inherit to create your own "fake" command."""

    def __repr__(self) -> str:
        self()

        return "\u200b"  # zero-width space

    @abc.abstractmethod
    def __call__(self) -> typing.Any:  # noqa: D102, ANN401
        pass


class _ClearConsole:
    """Clear command to... clear the terminal."""

    cursor_to_zero = "\r\x1b[H"
    clear_screen = "\x1b[J"

    def __repr__(self) -> str:
        is_pride_month = datetime.datetime.now().month == _PRIDE_MONTH_NUMBER

        sys.stdout.write(self.cursor_to_zero + self.clear_screen)
        sys.stdout.write(
            make_pill(
                get_pretty_python_version(),
                get_pretty_date(),
                *(
                    (
                        "\x1b[1m"
                        "\x1b[38;5;9mH"
                        "\x1b[38;5;1ma"
                        "\x1b[38;5;3mp"
                        "\x1b[38;5;11mp"
                        "\x1b[38;5;10my "
                        "\x1b[38;5;2mP"
                        "\x1b[38;5;14mr"
                        "\x1b[38;5;6mi"
                        "\x1b[38;5;12md"
                        "\x1b[38;5;4me "
                        "\x1b[38;5;5mM"
                        "\x1b[38;5;13mo"
                        "\x1b[38;5;15mn"
                        "\x1b[38;5;7mt"
                        "\x1b[38;5;8mh"
                        "\x1b[38;5;0m!"
                        "\x1b[22;39m",
                    )
                    if is_pride_month
                    else ()
                ),
                get_pretty_time(),
            ),
        )

        return ""

    def __call__(self) -> None:
        # clearing will leave a trailing newline so we move up beforehand for
        # consistency
        sys.stdout.write("\x1b[A")

        repr(self)


class _ExitConsole(FakeCommand):
    """Exit command (no parentheses needed!)."""

    def __call__(self, code: int = 0) -> typing.Never:
        raise SystemExit(code)


class _CalendarCommand(FakeCommand):
    """Command to display the calendar."""

    def __call__(self) -> None:
        print(_calendar.calendar(datetime.date.today().year, m=4))  # noqa: DTZ011


@dataclasses.dataclass(repr=False, frozen=True)
class ShellCommand(FakeCommand):
    """Run shell commands directly in your Python REPL. Do NOT abuse."""

    cmd: str

    def __call__(self) -> None:  # noqa: D102
        typing.cast(DynamicPS1, sys.ps1).status = subprocess.call([self.cmd])  # noqa: S603


#########
# UTILS #
#########


class _UtilsList(FakeCommand, list[AnyCallable]):
    """Utilitary to list and provide info on REPL utils."""

    __name__ = "utils"

    def __call__(self, function: AnyCallable | None = None) -> None:
        buffer = io.StringIO()

        if function is None:
            for element in self:
                buffer.write(f"\x1b[1;37m{element.__name__}\x1b[22;39m")

                if element.__doc__:
                    doc = textwrap.dedent(element.__doc__)
                    first_line = doc.strip().splitlines()[0]

                    if not first_line.endswith("."):
                        first_line += "..."

                    function_comment = (
                        f" \x1b[2m·\x1b[22m \x1b[3;92m{first_line}\x1b[23;39m"
                    )
                    buffer.write(f"{function_comment:>48}")

                buffer.write("\n")
        else:
            console = rich.console.Console(file=buffer, force_terminal=True)
            console.out(
                f"def {function.__name__}{inspect.signature(function, eval_str=True)}",
            )

            if function.__doc__:
                doc = function.__doc__.strip("\n")
                buffer.write(f"\x1b[92m{doc}\x1b[39m")

        builtins.print(buffer.getvalue())

    __annotations__ = __call__.__annotations__

    def register(
        self,
        function: AnyCallableT,
        *,
        with_name: str | None = None,
    ) -> AnyCallableT:
        _function = copy.copy(function)

        self.append(_function)

        if with_name:
            _function.__name__ = with_name

        self.sort(key=lambda function: function.__name__)

        return _function


utils: typing.Final = _UtilsList()
utils.append(utils)

literal: typing.Final = utils.register(ast.literal_eval, with_name="literal")


@utils.register
def concat(*strings: str, sep: str = " ") -> str:
    """Concatenate strings together giving a separator `sep` (default: " ")."""

    return sep.join(strings)


@utils.register
def clamp(number: int, _min: int = 0, _max: int = sys.maxsize) -> int:
    """Minmax an integer to a `min` (default: 0) and a `max` (default: sys.maxsize)."""

    return max(_min, min(_max, number))


@utils.register
def choose(*values: object) -> object:
    """
    Pick a random element among the values.

    Variadic equivalent of `random.choice()`.
    """

    return random.choice(values)


@utils.register
def compose(
    func1: collections.abc.Callable[[_T0], _T1],
    func2: collections.abc.Callable[[_T1], _T2],
) -> collections.abc.Callable[[_T0], _T2]:
    """
    Compose two functions together from left to right.

    >>> compose(str, len)(36)
    2
    """

    return lambda arg: func2(func1(arg))


@utils.register
def flipped(
    func: collections.abc.Callable[[_T0, _T1], _T2],
) -> collections.abc.Callable[[_T1, _T0], _T2]:
    """
    Return a version of the function with its arguments flipped.

    >>> flipped(operator.sub)(3, 5)
    2
    """

    return lambda arg1, arg2: func(arg2, arg1)


@utils.register
def raw_length(value: str) -> int:
    """Calculate the raw length of the string, i.e. its "displayed" length."""

    return compose(esclean, len)(value)


@utils.register
def float_equal(f1: float, f2: float) -> bool:
    """
    Estimate if two floats are equal using the machine's epsilon.

    >>> float_equal(0.1 + 0.2, 0.3)
    True
    """

    return 0 <= abs(f2 - f1) <= sys.float_info.epsilon


@utils.register
def irange(start: int, stop: int) -> range:
    """
    `range`, but with inclusive stop.

    >>> irange(0, 10)
    range(0, 11)
    """

    return range(start, stop + 1)


@utils.register
def magic(string: str) -> list[str]:
    """
    Convert a string into a list of the hexadecimal values of its characters.

    >>> magic("hello")
    ['0x68', '0x65', '0x6c', '0x6c', '0x6f']
    """

    return list(map(hex, string.encode("utf-8")))


@utils.register
def zeros_iter(n: int = -1, /) -> collections.abc.Generator[int, None, None]:
    """Return an iterator of `n` zeros."""

    i = 0

    while i < n or n < 0:
        yield 0
        i += 1


@utils.register
def zeros(n: int, /) -> list[int]:
    """
    Return a list of `n` zeros.

    >>> zeros(5)
    [0, 0, 0, 0, 0]
    >>> zeros(0)
    []
    >>> zeros(-2)
    *- ValueError: n must be non-negative -*
    """

    if n < 0:
        raise ValueError("n must be non-negative")

    return list(zeros_iter(n))


@utils.register
def bool_to_sign(value: bool, /) -> typing.Literal[1, -1]:  # noqa: FBT001
    """
    Return 1 if value is `True`, else -1.

    >>> bool_to_sign(True)
    1
    >>> bool_to_sign(False)
    -1
    """

    return 1 if value else -1


@utils.register
def parity_to_sign(value: int, /, *, inverted: bool = False) -> typing.Literal[1, -1]:
    """
    Return 1 if value is even, else -1. For consistency, 0 is considered even.

    If `inverted` is set to `True`, the result is negated.

    >>> parity_to_sign(2)
    1
    >>> parity_to_sign(-3)
    -1
    >>> parity_to_sign(0)
    1
    >>> parity_to_sign(40, inverted=True)
    -1
    """
    return bool_to_sign(value % 2 == inverted)


@utils.register
def fib(n: int) -> int:
    """The Fibonacci sequence over the integers."""

    if n < 0:
        return parity_to_sign(n, inverted=True) * fib(abs(n))
    if n <= 1:
        return n
    return fib(n - 2) + fib(n - 1)


@utils.register
def os_colorify(string: str) -> str:
    """Stylize the `string` with the color of the OS."""

    if sys.platform != "linux":
        raise OSError("this function can only run on Linux")

    info = platform.freedesktop_os_release()

    if "ANSI_COLOR" not in info:
        raise OSError("operating system does not have a color")

    return f"\x1b[{info['ANSI_COLOR']}m{string}\x1b[39m"


################################
# PRETTY PRINT OUTPUT & ERRORS #
################################


def get_pretty_python_version() -> str:
    """Render a colorful piece of text of the current Python version."""

    return "\x1b[1mPython\x1b[0m \x1b[91m{}.{}.{}\x1b[39m".format(*sys.version_info[:3])


def get_pretty_date() -> str:
    """Render a colorful piece of text of the current date."""

    return datetime.datetime.now().strftime("\x1b[1;35m%A %d %B %Y\x1b[22;39m")


def get_pretty_time() -> str:
    """Render a colorful piece of text of the current time."""

    return datetime.datetime.now().strftime("\x1b[1;34m%H:%M\x1b[22;39m")


@utils.register
def esclean(string: str) -> str:
    """
    Clean the string from escape sequences.

    Regex shamelessly stolen from a tutorialspoint post:
    https://www.tutorialspoint.com/How-can-I-remove-the-ANSI-escape-sequences-from-a-string-in-python
    """

    return re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]").sub("", string)


@utils.register
def make_pill(main: str, *others: str) -> str:
    """
    Render a contiguous array of blocks containing text.

    TODO: fix the case where the pill ends up being bigger than the available
    terminal width

    >>> make_pill("hello", "world")
    ╭───────┬───────╮
    │ hello │ world │
    ╰───────┴───────╯
    """

    buffer = io.StringIO()
    main_len = raw_length(main)
    others_len = [raw_length(other) for other in others]

    buffer.write("╭" + "─" * (main_len + 2))

    for length in others_len:
        buffer.write("┬" + "─" * (length + 2))

    buffer.write("╮\n")
    buffer.write("│ " + " │ ".join((main, *others)) + " │\n")
    buffer.write("╰" + "─" * (main_len + 2))

    for length in others_len:
        buffer.write("┴" + "─" * (length + 2))

    buffer.write("╯\n")

    return buffer.getvalue()


clear = utils.register(_ClearConsole(), with_name="clear")
exit = utils.register(_ExitConsole(), with_name="exit")  # noqa: A001
neofetch = ShellCommand("neofetch")
calendar = _CalendarCommand()


class DynamicPrompt(abc.ABC):
    """Prompt that dynamically adapts according to the state of the REPL."""

    def __init__(self) -> None:
        self.status = 0
        self.index = 0

    def __repr__(self) -> str:
        return self.prompt()

    @abc.abstractmethod
    def prompt(self) -> str:
        """Return the dynamic prompt."""


class DynamicPS1(DynamicPrompt):
    """Main prompt that dynamically adapts according to the state of the REPL."""

    def prompt(self) -> str:  # noqa: D102
        self.index += 1
        color = _c["LightRed"] if self.status else _c["LightGreen"]

        return "\x1b[m\x1b[2m[\x1b[22;35m%d\x1b[2;39m]\x1b[22m %sλ%s " % (
            self.index,
            color,
            _c["Normal"],
        )


class DynamicPS2(DynamicPrompt):
    """\
    Secondary prompt string that dynamically adapts according to what is
    happening in the REPL.
    """

    def prompt(self) -> str:  # noqa: D102
        self.index += 1
        typing.cast(DynamicPrompt, sys.ps1).index -= 1

        return f"\x1b[m{' ' * (raw_length(repr(sys.ps1)) - 2)}%s┊%s " % (
            _c["Purple"],
            _c["Normal"],
        )


# Make the prompts colorful.
sys.ps1 = DynamicPS1()
sys.ps2 = DynamicPS2()


# Enable pretty printing for STDOUT


def pretty_display_hook(value: object) -> None:
    """Custom display hook to make printed values colorful and pretty."""

    typing.cast(DynamicPS1, sys.ps1).status = 0

    if value is not None:
        __builtins__._ = value
        print(value)


sys.displayhook = pretty_display_hook

# Make errors and tracebacks stand out a bit more.


def pretty_except_hook(
    exc_type: type[BaseException],
    exc_value: BaseException,
    exc_tb: types.TracebackType | None,
) -> None:
    """Custom exception hook to make them look colorful and pretty."""

    typing.cast(DynamicPS1, sys.ps1).status = 1
    sys.stderr.write(_c["Yellow"])

    traceback.print_exception(exc_type, exc_value, exc_tb)

    sys.stderr.write(_c["Normal"])

    # NOTE: There is a bug (?) in Python 3, where a trailing color marker that's
    # written to STDERR or STDOUT by itself does not color the subsequent lines.
    # We work around this by manually calling ``flush`` afterwards.
    sys.stderr.flush()


# Python 3.13 introduces pretty exceptions so we don't need our custom excepthook
if sys.version_info < (3, 13):
    sys.excepthook = pretty_except_hook


_setrecursionlimit: typing.Final = sys.setrecursionlimit


def recursion_limit_setter_safe(limit: int, /) -> None:
    """\
    Set the recursion limit of the Python interpreter.

    This function overrides the built-in to prevent basic footguns.
    """

    if limit > __MAX_RECURSION_LIMIT:
        raise ValueError(
            "cannot set recursion limit to higher than 32768 "
            "\x1b[35m[prevent-large-recursion]\x1b[39m",
        )

    if limit < __MIN_RECURSION_LIMIT:
        raise ValueError(
            "cannot set recursion limit to lower than 2048 "
            "\x1b[35m[prevent-small-recursion]\x1b[39m",
        )

    _setrecursionlimit(limit)


# Prevents a footgun ^^
sys.setrecursionlimit = recursion_limit_setter_safe


# Funny things happen here vvv


@utils.register
def hex_encode(string: str) -> str:
    """
    Encode the string as the hexadecimal representation of its bytes.

    >>> hex_encode("hello")
    68656C6C6F
    """

    return f"{int.from_bytes(string.encode()):X}"


@utils.register
def hex_decode(encoded: str, length: int) -> str:
    """Decode the hexadecimal representation of a string's bytes into the \
    original object.

    >>> hex_decode("68656C6C6F", 5)
    hello
    """

    return int(encoded, 16).to_bytes(length).replace(b"\x00", b"").decode()


@utils.register
def encode_funcname(function: collections.abc.Callable[..., typing.Any]) -> str:
    """Encode a function as an unique identifier that can be used to retrieve \
    the function object later.

    >>> encode_funcname(abs)
    _6275696C74696E73_616273
    """

    return "_" + hex_encode(function.__module__) + "_" + hex_encode(function.__name__)


@utils.register
def decode_funcname(byteid: str) -> collections.abc.Callable[..., typing.Any]:
    """Decode the identifier generated by `encode_funcname` and retreive, if \
    it exists, the original function object.

    >>> decode_funcname("_6275696C74696E73_616273")
    <built-in function abs>
    """

    _, module_id, func_id = byteid.split("_")

    module_name = hex_decode(module_id, 64)
    func_name = hex_decode(func_id, 64)

    try:
        return getattr(importlib.import_module(module_name), func_name)
    except Exception:  # noqa: BLE001
        raise ValueError("invalid func byte id") from None


src: typing.Final = utils.register(inspect.getsource, with_name="src")


def help(request: object = None, /) -> None:  # noqa: A001, D103
    builtins.help(request)
    sys.stdout.write("\x1b[?1049h")
    builtins.print(clear)


@atexit.register
def _exit_alt_buffer() -> None:  # pyright: ignore[reportUnusedFunction]
    """\
    Disable the alternative screen buffer at exit.
    """

    import sys  # for IPython

    sys.stdout.write("\x1b[?1049l")
    sys.stdout.flush()
    sys.stdout.write(
        "\x1b[32m♡ \x1b[1mThank you for using the Python REPL powered with "
        "qexat's .pythonrc!\x1b[22;39m\n",
    )


__has_initialized = False


def init(_: list[str]) -> None:
    """
    The default init function of the REPL.

    Tasks
    -----
    * Print a launching message
    * Enable the fault handler
    * Set the recursion limit to a higher, but safe limit
    * Switch to the alternative screen buffer
    * Clear the screen

    If the REPL is already initialized, do nothing.
    """

    if __has_initialized:
        return

    sys.stdout.write(LAUNCHING_MESSAGE)

    # Sometimes I do some fuckery in the REPL so we should get ourselves covered
    faulthandler.enable()
    # 1000 is too low...
    sys.setrecursionlimit(__MAX_RECURSION_LIMIT)

    # Enter the alternative screen buffer
    # We do NOT use `print` because `rich` strips the escape sequence somehow
    sys.stdout.write("\x1b[?1049h")
    sys.stdout.flush()

    # We clear the buffer so it looks like a new window
    builtins.print(clear)


def main(
    init_func: collections.abc.Callable[[list[str]], None],
    *,
    args: list[str],
) -> None:
    """
    Entry point of the REPL.

    Parameters
    ----------
    init_func : (list[str]) -> None
        A function run at the initialization of the REPL.
    args : list[str]
        The CLI arguments passed to the REPL.

    """

    global __has_initialized  # noqa: PLW0603

    if __has_initialized:
        builtins.print(
            "\x1b[1;93mWARNING: \x1b[39mREPL is already initialized\x1b[22m",
            file=sys.stderr,
        )

    init_func(args)
    __has_initialized = True


if __name__ == "__main__":
    main(init, args=sys.argv[1:])