bitflag_visualize.py

#!/usr/bin/env python3
"""Visualize an integer as bits and optionally label single-bit masks.

Usage: python3.14 bitflag_visualize.py NUMBER [--flags JSON5_OR_FILE]
Requires the ``json5`` package. NUMBER may be decimal, hexadecimal, or binary.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import cast

import json5


class Colors:
    """ANSI colors used by the terminal table."""
    HEADER = "\033[95m"
    BLUE = "\033[94m"
    CYAN = "\033[96m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    RED = "\033[91m"
    GREY = "\033[90m"
    BOLD = "\033[1m"
    ENDC = "\033[0m"


def parse_int(value: object) -> int:
    val_str = str(value).strip().replace("_", "")
    try:
        if val_str.lower().startswith("0x"):
            return int(val_str, 16)
        elif val_str.lower().startswith("0b"):
            return int(val_str, 2)
        else:
            return int(val_str)
    except ValueError:
        raise ValueError(f"Could not parse number: {value!r}") from None


def load_flags_from_masks(flags_arg: str | None) -> dict[int, str]:
    if not flags_arg:
        return {}

    data_str = ""
    if Path(flags_arg).is_file():
        try:
            with open(flags_arg, encoding="utf-8") as f:
                data_str = f.read()
        except Exception as e:
            raise ValueError(f"Could not read flags file: {e}") from e
    else:
        data_str = flags_arg

    try:
        raw_dict = cast(dict[str, object], json5.loads(data_str))
        bit_map: dict[int, str] = {}
        for name, mask_raw in raw_dict.items():
            if isinstance(mask_raw, str):
                mask_val = parse_int(mask_raw)
            elif isinstance(mask_raw, int):
                mask_val = mask_raw
            else:
                continue

            if mask_val > 0 and (mask_val & (mask_val - 1) == 0):
                bit_index = mask_val.bit_length() - 1
                bit_map[bit_index] = name
        return bit_map

    except Exception as e:
        raise ValueError(f"Could not parse JSON5 flags: {e}") from e


def print_table(value: int, flag_map: dict[int, str]) -> None:
    # Determine the bit range.
    max_bit_val = value.bit_length()
    max_flag_pos = max(flag_map.keys()) if flag_map else 0
    total_bits = max(max_bit_val, max_flag_pos + 1)

    # Round the width up to a complete byte.
    width = ((total_bits + 7) // 8) * 8
    if width == 0:
        width = 8

    # Calculate the mask column width from the highest bit.
    max_mask_val = 1 << (width - 1)
    mask_col_len = len(f"0x{max_mask_val:X}")
    header_mask_title = "Mask (Hex)"
    col_width = max(mask_col_len, len(header_mask_title))

    total_table_len = 6 + 3 + col_width + 3 + 5 + 3 + 20

    print(f"\nAnalysis: {Colors.BOLD}{value}{Colors.ENDC}")
    print(f"HEX   : {Colors.CYAN}{hex(value)}{Colors.ENDC}")
    print(f"BIN   : {Colors.YELLOW}{bin(value)}{Colors.ENDC}")
    print("-" * total_table_len)

    print(
        f"{Colors.HEADER}{'Bit #':<6} | {header_mask_title:<{col_width}} | {'Val':<5} | {'Flag Name'}{Colors.ENDC}"
    )
    print("-" * total_table_len)

    for i in range(width - 1, -1, -1):
        mask = 1 << i
        is_set = (value & mask) != 0

        flag_name = flag_map.get(i, "")

        mask_str = f"0x{mask:X}"
        bit_idx_str = f"{i:<6}"
        val_display = " 1 " if is_set else " 0 "

        if is_set:
            val_colored = f"{Colors.GREEN}{val_display}{Colors.ENDC}"
            bit_idx_colored = f"{Colors.BOLD}{bit_idx_str}{Colors.ENDC}"
            mask_colored = f"{Colors.CYAN}{mask_str:<{col_width}}{Colors.ENDC}"
            name_colored = (
                f"{Colors.GREEN}{Colors.BOLD}{flag_name}{Colors.ENDC}" if flag_name else ""
            )
        else:
            val_colored = f"{Colors.GREY}{val_display}{Colors.ENDC}"
            bit_idx_colored = bit_idx_str
            mask_colored = f"{mask_str:<{col_width}}"
            name_colored = ""

        print(f"{bit_idx_colored} | {mask_colored} | {val_colored} | {name_colored}")

    print("-" * total_table_len + "\n")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Bit analyzer with JSON5 masks and automatic width"
    )
    parser.add_argument("number", help="number to analyze")
    parser.add_argument("--flags", "-f", help="JSON5 string or file containing masks")

    args = parser.parse_args(argv)

    try:
        target_value = parse_int(args.number)
        flag_map = load_flags_from_masks(args.flags)
    except ValueError as exc:
        parser.error(str(exc))

    print_table(target_value, flag_map)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())