Redesign/Approach to _colorize.py using Types

_colorize.py

"""
https://github.com/python/cpython/blob/main/Lib/_colorize.py

## Redesign/Approach to _colorize.py

1. Improve ergonomics of applying a theme from f"{t.BLUE} text {t.reset}" to f"{t.BLUE(' text ')}".
2. Leverage Types to improve readability. Specifically, ANSIColor to be a specific type.
3. NoColor vs Color is similar to None | int from a type standpoint.
4. There's a bit going on with `ThemeSection(Mapping[str,str])`. Maybe this can be streamlined.
5. `ThemeSection.no_colors()` can be simplified by #3
6. `ThemeSection.copy_with()` can be simplified by using `__replace__` and #3
7. `ThemeSection` has `__post_init__` with ` super().__setattr__('_name_to_value', name_to_value.__getitem__)`. This is not super easy to understand.


## Changes

At a high level, the main change is thinking in terms of types and leverage __replace__ on dataclass

1. ANSIColor to StrEnum
2. define Color(AnsiColor) layer with __call__
3. define __replace__ on ThemeSection and Theme. Use copy.replace to replace theme section, or color
4. Add demo() and display_demo() method for kicking the tires


Note: Backwards compatibility requirements might dramatically change the design.
"""

import copy
from enum import StrEnum
from typing import Protocol, ClassVar, Iterator, Self
from dataclasses import dataclass, Field, fields, field


class ANSIColors(StrEnum):
    RESET = "\x1b[0m"

    BLACK = "\x1b[30m"
    BLUE = "\x1b[34m"
    CYAN = "\x1b[36m"
    GREEN = "\x1b[32m"
    GREY = "\x1b[90m"
    MAGENTA = "\x1b[35m"
    RED = "\x1b[31m"
    WHITE = "\x1b[37m"  # more like LIGHT GRAY
    YELLOW = "\x1b[33m"

    BOLD = "\x1b[1m"
    BOLD_BLACK = "\x1b[1;30m"  # DARK GRAY
    BOLD_BLUE = "\x1b[1;34m"
    BOLD_CYAN = "\x1b[1;36m"
    BOLD_GREEN = "\x1b[1;32m"
    BOLD_MAGENTA = "\x1b[1;35m"
    BOLD_RED = "\x1b[1;31m"
    BOLD_WHITE = "\x1b[1;37m"  # actual WHITE
    BOLD_YELLOW = "\x1b[1;33m"

    # intense = like bold but without being bold
    INTENSE_BLACK = "\x1b[90m"
    INTENSE_BLUE = "\x1b[94m"
    INTENSE_CYAN = "\x1b[96m"
    INTENSE_GREEN = "\x1b[92m"
    INTENSE_MAGENTA = "\x1b[95m"
    INTENSE_RED = "\x1b[91m"
    INTENSE_WHITE = "\x1b[97m"
    INTENSE_YELLOW = "\x1b[93m"

    BACKGROUND_BLACK = "\x1b[40m"
    BACKGROUND_BLUE = "\x1b[44m"
    BACKGROUND_CYAN = "\x1b[46m"
    BACKGROUND_GREEN = "\x1b[42m"
    BACKGROUND_MAGENTA = "\x1b[45m"
    BACKGROUND_RED = "\x1b[41m"
    BACKGROUND_WHITE = "\x1b[47m"
    BACKGROUND_YELLOW = "\x1b[43m"

    INTENSE_BACKGROUND_BLACK = "\x1b[100m"
    INTENSE_BACKGROUND_BLUE = "\x1b[104m"
    INTENSE_BACKGROUND_CYAN = "\x1b[106m"
    INTENSE_BACKGROUND_GREEN = "\x1b[102m"
    INTENSE_BACKGROUND_MAGENTA = "\x1b[105m"
    INTENSE_BACKGROUND_RED = "\x1b[101m"
    INTENSE_BACKGROUND_WHITE = "\x1b[107m"
    INTENSE_BACKGROUND_YELLOW = "\x1b[103m"


class Color:
    def __init__(self, ansi_escape: ANSIColors):
        self._ansi_escape = ansi_escape

    def __repr__(self) -> str:
        # escape the value to show the raw value
        return f"<Color \\{self._ansi_escape} >"

    @property
    def value(self) -> str:
        return self._ansi_escape.value

    def __call__(self, text: str) -> str:
        return f"{self._ansi_escape.value}{text}{ANSIColors.RESET.value}"


class NoColor:

    def __init__(self) -> None:
        # this should be a singleton
        pass

    def __repr__(self) -> str:
        return "<NoColor>"

    @property
    def value(self) -> str:
        return ""

    def __call__(self, text: str) -> str:
        return text


ColorAble =  Color | NoColor

# alternatively this could be done using Structural Typing as
class ColorAble2(Protocol):
    value: str

    def __call__(self, text: str) -> str:
        ...


# Ideally this would be in the stdlib
# https://github.com/python/cpython/issues/102699
class DataclassAble[T](Protocol):
    __dataclass_fields__: ClassVar[dict[str, Field[T]]]



class ThemeSection(DataclassAble[ColorAble]):
    # this is abstract due to issues with mixin to
    # dataclasses creates type confusion
    reset: ColorAble

    def display_demo(self, text: str = "example text", pad: int = 2) -> Iterator[str]:
        yield f"Theme {self.__class__.__name__}"
        max_length = max(len(x.name) for x in fields(self)) + pad
        for f in fields(self):
            yield f"{f.name:<{max_length}}{getattr(self, f.name)(text)}"

    @classmethod
    def no_colors(cls) -> Self:
        n = NoColor()
        return cls(**{name:n for name in cls.__dataclass_fields__})



@dataclass(frozen=True, kw_only=True)
class Argparse(ThemeSection):
    usage: ColorAble = Color(ANSIColors.BOLD_BLUE)
    prog: ColorAble = Color(ANSIColors.BOLD_MAGENTA)
    prog_extra: ColorAble = Color(ANSIColors.MAGENTA)
    heading: ColorAble = Color(ANSIColors.BOLD_BLUE)
    summary_long_option: ColorAble = Color(ANSIColors.CYAN)
    summary_short_option: ColorAble = Color(ANSIColors.GREEN)
    summary_label: ColorAble = Color(ANSIColors.YELLOW)
    summary_action: ColorAble = Color(ANSIColors.GREEN)
    long_option: ColorAble = Color(ANSIColors.BOLD_CYAN)
    short_option: ColorAble = Color(ANSIColors.BOLD_GREEN)
    label: ColorAble = Color(ANSIColors.BOLD_YELLOW)
    action: ColorAble = Color(ANSIColors.BOLD_GREEN)
    reset: ColorAble = Color(ANSIColors.RESET)


@dataclass(frozen=True, kw_only=True)
class Difflib(ThemeSection):
    """A 'git diff'-like theme for `difflib.unified_diff`."""
    added: ColorAble = Color(ANSIColors.GREEN)
    context: ColorAble = Color(ANSIColors.RESET)  # context lines
    header: ColorAble = Color(ANSIColors.BOLD)  # eg "---" and "+++" lines
    hunk: ColorAble = Color(ANSIColors.CYAN)  # the "@@" lines
    removed: ColorAble = Color(ANSIColors.RED)
    reset: ColorAble = Color(ANSIColors.RESET)


@dataclass(frozen=True, kw_only=True)
class Syntax(ThemeSection):
    prompt: ColorAble = Color(ANSIColors.BOLD_MAGENTA)
    keyword: ColorAble = Color(ANSIColors.BOLD_BLUE)
    keyword_constant: ColorAble = Color(ANSIColors.BOLD_BLUE)
    builtin: ColorAble = Color(ANSIColors.CYAN)
    comment: ColorAble = Color(ANSIColors.RED)
    string: ColorAble = Color(ANSIColors.GREEN)
    number: ColorAble = Color(ANSIColors.YELLOW)
    op: ColorAble = Color(ANSIColors.RESET)
    definition: ColorAble = Color(ANSIColors.BOLD)
    soft_keyword: ColorAble = Color(ANSIColors.BOLD_BLUE)
    reset: ColorAble = Color(ANSIColors.RESET)


@dataclass(frozen=True, kw_only=True)
class Traceback(ThemeSection):
    type: ColorAble = Color(ANSIColors.BOLD_MAGENTA)
    message: ColorAble = Color(ANSIColors.MAGENTA)
    filename: ColorAble = Color(ANSIColors.MAGENTA)
    line_no: ColorAble = Color(ANSIColors.MAGENTA)
    frame: ColorAble = Color(ANSIColors.MAGENTA)
    error_highlight: ColorAble = Color(ANSIColors.BOLD_RED)
    error_range: ColorAble = Color(ANSIColors.RED)
    reset: ColorAble = Color(ANSIColors.RESET)


@dataclass(frozen=True, kw_only=True)
class Unittest(ThemeSection):
    passed: ColorAble = Color(ANSIColors.GREEN)
    warn: ColorAble = Color(ANSIColors.YELLOW)
    fail: ColorAble = Color(ANSIColors.RED)
    fail_info: ColorAble = Color(ANSIColors.BOLD_RED)
    reset: ColorAble = Color(ANSIColors.RESET)


@dataclass(frozen=True, kw_only=True)
class Theme:
    """A suite of themes for all sections of Python.

    When adding a new one, remember to also modify `copy_with` and `no_colors`
    below.
    """
    argparse: Argparse = field(default_factory=Argparse)
    difflib: Difflib = field(default_factory=Difflib)
    syntax: Syntax = field(default_factory=Syntax)
    traceback: Traceback = field(default_factory=Traceback)
    unittest: Unittest = field(default_factory=Unittest)

    @classmethod
    def no_colors(cls) -> Self:
        """Return a new Theme where colors in all sections are empty strings.

        This allows writing user code as if colors are always used. The color
        fields will be ANSI color code strings when colorization is desired
        and possible, and empty strings otherwise.
        """
        return cls(
            argparse=Argparse.no_colors(),
            difflib=Difflib.no_colors(),
            syntax=Syntax.no_colors(),
            traceback=Traceback.no_colors(),
            unittest=Unittest.no_colors(),
        )

    def display_demo(self, text: str = "example text") -> Iterator[str]:
        """Display the demo text."""
        for ts in (self.argparse, self.difflib, self.syntax, self.traceback, self.unittest):
            yield from ts.display_demo(text)



def demo() -> Theme:
    print(repr(ANSIColors.BOLD_BLUE))
    a = Argparse()
    print(a.heading("This is the Heading"))

    b = copy.replace(a, heading=NoColor(), prog=Color(ANSIColors.BOLD_BLUE))
    print(b.heading("This is the Heading"))

    t1 = Theme()
    t2 = copy.replace(t1, argparse=b)

    for t in (t1, t2):
        print("\n".join(t.argparse.display_demo()))

    return t