unfuck_python_colors.py

#!/usr/bin/env python

# Unfucks Python 3.13+ REPL colors for Windows.
# See context: https://github.com/python/cpython/issues/125609

from __future__ import annotations

import os
import site
import sys
import argparse
import importlib.util

# Default colors that might be more readable on Windows terminals
DEFAULT_COLOR_OVERRIDES = {
    'MAGENTA': '\\x1b[0;38;5;213m',      # Brighter magenta
    'BOLD_MAGENTA': '\\x1b[1;38;5;213m', # Brighter bold magenta
    'RED': '\\x1b[0;38;5;203m',          # Brighter red
    'BOLD_RED': '\\x1b[1;38;5;203m'      # Brighter bold red
}

SITECUSTOMIZE_TEMPLATE = """# Generated by unfuck_python_colors.py
# This file customizes the colors used in the Python 3.13+ colorized REPL.

import importlib.util

if importlib.util.find_spec("_colorize"):
    import _colorize
    
    # Color customizations
{color_customizations}
    
"""

def get_site_packages() -> str:
    """Get the site-packages directory."""
    return site.getsitepackages()[0]

def check_if_module_exists(module_name: str) -> bool:
    """Check if a Python module exists."""
    return importlib.util.find_spec(module_name) is not None

def confirm(prompt: str) -> bool:
    """Simple confirmation prompt."""
    response = input(f'{prompt} (y/n): ').lower().strip()
    return response in ('y', 'yes')

def install(magenta: str = DEFAULT_COLOR_OVERRIDES['MAGENTA'], 
            bold_magenta: str = DEFAULT_COLOR_OVERRIDES['BOLD_MAGENTA'],
            red: str = DEFAULT_COLOR_OVERRIDES['RED'], 
            bold_red: str = DEFAULT_COLOR_OVERRIDES['BOLD_RED'],
            force: bool = False) -> None:
    """Install sitecustomize.py to customize Python REPL colors."""
    # Check Python version
    if sys.version_info < (3, 13):
        print('WARNING: This script is intended for Python 3.13+ which has the colorized REPL.')
        if not confirm('Continue anyway?'):
            return

    # Check if _colorize module exists
    if not check_if_module_exists('_colorize'):
        print('WARNING: The _colorize module was not found. Are you using Python 3.13+?')
        if not confirm('Continue anyway?'):
            return
    
    # Get the user site-packages directory
    site_packages = get_site_packages()
    
    # Ensure the directory exists
    os.makedirs(site_packages, exist_ok=True)
    
    # Path to sitecustomize.py
    sitecustomize_path = os.path.join(site_packages, 'sitecustomize.py')
    
    # Check if file already exists
    if os.path.exists(sitecustomize_path) and not force:
        print(f'A sitecustomize.py file already exists at {sitecustomize_path}')
        if not confirm('Overwrite it?'):
            return
    
    # Build color customizations string
    color_overrides = {
        'MAGENTA': magenta,
        'BOLD_MAGENTA': bold_magenta,
        'RED': red,
        'BOLD_RED': bold_red
    }
    
    color_customizations = '\n'.join([
        f"    _colorize.ANSIColors.{color} = '{value}'" 
        for color, value in color_overrides.items()
    ])
    
    # Generate the content
    content = SITECUSTOMIZE_TEMPLATE.format(color_customizations=color_customizations)
    
    # Write the file
    with open(sitecustomize_path, 'w') as f:
        f.write(content)
    
    print(f'Successfully created {sitecustomize_path}')
    print('Restart your Python interpreter for changes to take effect.')
    
    if not site.ENABLE_USER_SITE:
        print('\nRemember: In virtual environments, you may need to:')
        print('1. Create the venv with --system-site-packages, or')
        print('2. Set PYTHONPATH to include the user site-packages directory')

def info() -> None:
    """Show information about the current Python environment."""
    print(f'Python version: {sys.version}')
    print(f'Site-packages: {get_site_packages()}')
    print(f'_colorize module available: {check_if_module_exists("_colorize")}')
    
    # Show sys.path
    print('\nPython path (sys.path):')
    for i, path in enumerate(sys.path, 1):
        print(f'{i}. {path}')

def list_colors() -> None:
    """List all the available ANSI color codes."""
    if not check_if_module_exists('_colorize'):
        print('The _colorize module is not available. Are you using Python 3.13+?')
        return
    
    import _colorize
    print('Available colors in _colorize.ANSIColors:')
    for attr in dir(_colorize.ANSIColors):
        if not attr.startswith('_'):
            value = getattr(_colorize.ANSIColors, attr)
            print(f"{attr} = '{value}'")

def main() -> None:
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description='Tool to customize Python 3.13+ REPL colors for better readability.'
    )
    subparsers = parser.add_subparsers(dest='command', help='Command to run')
    
    # Install command
    install_parser = subparsers.add_parser('install', help='Install sitecustomize.py')
    install_parser.add_argument('--magenta', default=DEFAULT_COLOR_OVERRIDES['MAGENTA'],
                               help='ANSI color code for MAGENTA')
    install_parser.add_argument('--bold-magenta', default=DEFAULT_COLOR_OVERRIDES['BOLD_MAGENTA'],
                               help='ANSI color code for BOLD_MAGENTA')
    install_parser.add_argument('--red', default=DEFAULT_COLOR_OVERRIDES['RED'],
                               help='ANSI color code for RED')
    install_parser.add_argument('--bold-red', default=DEFAULT_COLOR_OVERRIDES['BOLD_RED'],
                               help='ANSI color code for BOLD_RED')
    install_parser.add_argument('--force', action='store_true', 
                               help='Overwrite existing sitecustomize.py file')
    
    # Info command
    subparsers.add_parser('info', help='Show information about the current Python environment')
    
    # List colors command
    subparsers.add_parser('list-colors', help='List all the available ANSI color codes')
    
    args = parser.parse_args()
    
    print(sys.executable)
    
    if args.command == 'install':
        install(
            magenta=args.magenta,
            bold_magenta=args.bold_magenta,
            red=args.red,
            bold_red=args.bold_red,
            force=args.force
        )
    elif args.command == 'info':
        info()
    elif args.command == 'list-colors':
        list_colors()
    else:
        parser.print_help()

if __name__ == '__main__':
    main()