.pythonstartup

"""
Python Startup Script – Interactive Interpreter Package Loader
==============================================================
Automatically imports common packages into the interactive session when Python
is launched. Missing packages are installed via pip before being re-imported.

Windows setup
-------------
Set the PYTHONSTARTUP environment variable to point to this file:
    setx PYTHONSTARTUP "C:\\path\\to\\pythonrc.py"
Then reopen your terminal and run: python

Customisation
-------------
Add or remove entries from PACKAGE_SHORTCUTS:
    { 'package_name': 'shortcut', ... }
"""

import importlib
import importlib.util
import os
import subprocess
import sys
from types import ModuleType
from typing import Optional


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

PACKAGE_SHORTCUTS: dict[str, str] = {
    'os': 'os',
    'sys': 'sys',
    'numpy': 'np',
    'scipy': 'sp',
    'pandas': 'pd',
    'sympy': 'sym',
    'matplotlib.pyplot': 'plt',
    'pathlib': 'pathlib',
}

# ---------------------------------------------------------------------------
# ANSI – enable VT processing on Windows, then define colour/style codes
# ---------------------------------------------------------------------------


def _enable_windows_ansi() -> bool:
    """Enable VT-100 escape processing in the Windows console. No-op elsewhere."""
    if os.name != 'nt':
        return True
    try:
        import ctypes
        import ctypes.wintypes

        kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]
        handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
        mode = ctypes.wintypes.DWORD()
        kernel32.GetConsoleMode(handle, ctypes.byref(mode))
        kernel32.SetConsoleMode(handle, mode.value | 0x0004)  # ENABLE_VIRTUAL_TERMINAL_PROCESSING
        return True
    except Exception:
        return False


_USE_COLOUR = _enable_windows_ansi() and sys.stdout.isatty()


def _c(code: str) -> str:
    return code if _USE_COLOUR else ''


# Styles
RESET = _c('\033[0m')
BOLD = _c('\033[1m')
DIM = _c('\033[2m')
# Colours
CYAN = _c('\033[96m')
GREEN = _c('\033[92m')
YELLOW = _c('\033[93m')
RED = _c('\033[91m')
WHITE = _c('\033[97m')
GREY = _c('\033[90m')


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def cl() -> None:
    """Clear the console (cross-platform)."""
    os.system('cls' if os.name == 'nt' else 'clear')


def _target_namespace() -> dict:
    """
    Return the namespace visible to the user in the REPL (__main__).

    PYTHONSTARTUP runs inside __main__, so shortcuts must land there.
    Using globals() inside a helper would silently populate an unreachable
    module-level namespace instead.
    """
    return sys.modules['__main__'].__dict__


def _get_env_label() -> str:
    """Return a short label for the active Python environment."""
    venv = os.environ.get('VIRTUAL_ENV') or os.environ.get('CONDA_DEFAULT_ENV')
    if venv:
        return os.path.basename(venv)
    return 'system'


def _print_header() -> None:
    """Print the startup banner with Python version and environment info."""
    version = '{}.{}.{}'.format(*sys.version_info[:3])
    env = _get_env_label()
    content = f'  Python {version}   env: {env}  '
    width = len(content)
    print()
    print(f'{CYAN}{BOLD} ╭{"─" * width}╮{RESET}')
    print(f'{CYAN}{BOLD} │{WHITE}{content}{CYAN}│{RESET}')
    print(f'{CYAN}{BOLD} ╰{"─" * width}╯{RESET}')
    print()


def _print_table_header(col_pkg: int, col_alias: int) -> None:
    """Print the column header row and separator."""
    total = col_pkg + col_alias + 24
    print(f'   {BOLD}{WHITE}{"Package":<{col_pkg}}  {"Alias":<{col_alias}}  Status{RESET}')
    print(f'   {DIM}{"─" * total}{RESET}')


def _print_table_footer(col_pkg: int, col_alias: int, n_ok: int, n_fail: int, n_cached: int) -> None:
    """Print the separator and summary line."""
    total = col_pkg + col_alias + 24
    print(f'   {DIM}{"─" * total}{RESET}')
    parts = []
    if n_ok:
        parts.append(f'{GREEN}{BOLD}{n_ok} loaded{RESET}')
    if n_cached:
        parts.append(f'{GREY}{n_cached} cached{RESET}')
    if n_fail:
        parts.append(f'{RED}{BOLD}{n_fail} failed{RESET}')
    join = f'  {DIM}·{RESET}  '
    print(f'   {join.join(parts) if parts else f"{GREY}nothing to load{RESET}"}')
    print()


def _fmt_time(ms: float) -> str:
    """Format a duration: milliseconds below 1 s, seconds above."""
    if ms < 1000:
        return f'{ms:>5.0f} ms'
    return f'{ms / 1000:>5.2f}  s'


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------


def _is_installed(package_name: str) -> bool:
    """Return True if *package_name* can be found without importing it."""
    return importlib.util.find_spec(package_name) is not None


def _import_package(shortcut: str, package_name: str) -> Optional[ModuleType]:
    """Import *package_name* and bind it to *shortcut* in the REPL namespace."""
    try:
        module = importlib.import_module(package_name)
        _target_namespace()[shortcut] = module
        return module
    except ImportError:
        return None


def _install_package(package_name: str) -> bool:
    """
    Install *package_name* via pip using the current interpreter.
    sys.executable guarantees the correct env (venv-aware).
    Returns True on success.
    """
    try:
        subprocess.check_call(
            [sys.executable, '-m', 'pip', 'install', package_name, '-q'],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return True
    except subprocess.CalledProcessError:
        return False


def _load_packages(shortcuts: dict[str, str]) -> None:
    """Import every package in *shortcuts*, installing any that are missing."""
    import time

    _print_header()

    col_pkg = max((len(p) for p in shortcuts), default=10)
    col_alias = max((len(a) for a in shortcuts.values()), default=5)

    _print_table_header(col_pkg, col_alias)

    n_ok = n_fail = n_cached = 0

    for package_name, shortcut in shortcuts.items():
        pkg_col = f'{package_name:<{col_pkg}}'
        alias_col = f'{shortcut:<{col_alias}}'
        prefix = f'   {DIM}  {pkg_col}  {alias_col}{RESET}  '

        # ── Already in session namespace ──────────────────────────────────
        if shortcut in _target_namespace():
            print(f'{prefix}{GREY}· cached{RESET}')
            n_cached += 1
            continue

        # ── Attempt import ────────────────────────────────────────────────
        t0 = time.perf_counter()
        module = _import_package(shortcut, package_name)
        ms = (time.perf_counter() - t0) * 1000

        if module is not None:
            print(f'{prefix}{GREEN}✔ loaded   {DIM}{_fmt_time(ms)}{RESET}')
            n_ok += 1
            continue

        # ── Package absent → try installing ───────────────────────────────
        if not _is_installed(package_name):
            print(f'{prefix}{YELLOW}↓ installing …{RESET}', end='\r', flush=True)
            ok = _install_package(package_name)
            if ok:
                t0 = time.perf_counter()
                module = _import_package(shortcut, package_name)
                ms = (time.perf_counter() - t0) * 1000
                if module is not None:
                    print(f'{prefix}{YELLOW}⬇ installed  {DIM}{_fmt_time(ms)}{RESET}')
                    n_ok += 1
                    continue
            print(f'{prefix}{RED}✘ install failed   → pip install {package_name}{RESET}')
        else:
            print(f'{prefix}{RED}✘ import error     (installed but broken){RESET}')

        n_fail += 1

    _print_table_footer(col_pkg, col_alias, n_ok, n_fail, n_cached)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

_load_packages(PACKAGE_SHORTCUTS)

# Clean up: remove every name this script introduced from the REPL namespace
# so the user's session is not polluted with internal helpers.
# Cache the dict reference first — the loop would otherwise call
# _target_namespace() after that very name has already been popped.
_ns = _target_namespace()
for _name in (
    # public helper kept in session
    'cl',
    # internal functions
    '_load_packages',
    '_import_package',
    '_install_package',
    '_is_installed',
    '_target_namespace',
    '_get_env_label',
    '_enable_windows_ansi',
    '_print_header',
    '_print_table_header',
    '_print_table_footer',
    '_fmt_time',
    # constants & style codes
    '_USE_COLOUR',
    '_c',
    'RESET',
    'BOLD',
    'DIM',
    'CYAN',
    'GREEN',
    'YELLOW',
    'RED',
    'WHITE',
    'GREY',
    # stdlib imports surfaced at module level
    'importlib',
    'os',
    'subprocess',
    'ModuleType',
    'Optional',
    # config
    'PACKAGE_SHORTCUTS',
):
    _ns.pop(_name, None)
del _ns, _name