text Calendar in python with ANSI colors #py

cal.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import calendar
import datetime

# import sys

# ANSI color codes for terminal output
BLUE = "\033[94m"
RED = "\033[91m"
BOLD = "\033[1m"
HIGHLIGHT = "\033[46;30m"
DIM_WHITE = "\033[90m"
HOLIDAY_COLOR = "\033[4;95m"  # Underlined Magenta
RESET = "\033[0m"


def get_us_holidays(year):
    """Returns a dictionary of {(month, day): "Holiday Name"} for the US."""
    holidays = {
        (1, 1): "New Year's",
        (6, 19): "Juneteenth",
        (7, 4): "July 4th",
        # (11, 11): "Veterans Day",
        (12, 25): "Christmas",
    }

    # Helper to find the N-th occurrence of a weekday in a month
    def nth_weekday(year, month, weekday, n):
        # weekday 0=Mon, 3=Thu
        cal = calendar.monthcalendar(year, month)
        days = [week[weekday] for week in cal if week[weekday] != 0]
        return days[min(n - 1, len(days) - 1)]

    # Floating holidays
    holidays[(1, nth_weekday(year, 1, 0, 3))] = "MLK Day"
    holidays[(2, nth_weekday(year, 2, 0, 3))] = "Presidents Day"
    holidays[(9, nth_weekday(year, 9, 0, 1))] = "Labor Day"
    # holidays[(10, nth_weekday(year, 10, 0, 2))] = "Columbus Day"
    holidays[(11, nth_weekday(year, 11, 3, 4))] = "Thanksgiving"

    # Memorial Day is the LAST Monday in May
    may_cal = calendar.monthcalendar(year, 5)
    last_monday = may_cal[-1][0] if may_cal[-1][0] != 0 else may_cal[-2][0]
    holidays[(5, last_monday)] = "Memorial Day"

    return holidays


def format_month_list(year, month, today=None):
    """Returns a list of strings representing the calendar month."""
    cal = calendar.TextCalendar(calendar.SUNDAY)
    us_holidays = get_us_holidays(year)
    month_str = cal.formatmonth(year, month).splitlines()

    # Custom headers: Blue for month/days, Red for Su/Sa
    # header = f"{BLUE}{month_str[0].center(20)}{RESET}"

    raw_header = month_str[0].strip()
    header = f"{BLUE}{BOLD}{raw_header.center(20)}{RESET}"

    days_header = f"{RED}Su{RESET} {BLUE}Mo Tu We Th Fr{RESET} {RED}Sa{RESET}"

    formatted_weeks = [header, days_header]

    # Process weeks to highlight 'today'
    month_days = cal.monthdayscalendar(year, month)
    for week in month_days:
        week_strs = []
        for i, day in enumerate(week):
            if day == 0:
                week_strs.append("  ")
            elif (
                today
                and today.year == year
                and today.month == month
                and today.day == day
            ):
                week_strs.append(f"{HIGHLIGHT}{day:2}{RESET}")
            elif (month, day) in us_holidays:  # Add this check
                week_strs.append(f"{HOLIDAY_COLOR}{day:2}{RESET}")
            elif i == 0 or i == 6:  # Sun or Sat
                week_strs.append(f"{DIM_WHITE}{day:2}{RESET}")
            else:
                week_strs.append(f"{day:2}")
        formatted_weeks.append(" ".join(week_strs))

    # Ensure all months have 8 lines for consistent grid alignment
    while len(formatted_weeks) < 8:
        formatted_weeks.append(" " * 20)

    return formatted_weeks


def print_months_side_by_side(months_data):
    """Prints months side-by-side."""
    # Group months into rows of 3
    for i in range(0, len(months_data), 3):
        chunk = months_data[i : i + 3]
        for line_idx in range(8):
            line_parts = []
            for m in chunk:
                line_parts.append(m[line_idx].ljust(20))
            print("    ".join(line_parts))
        print()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "vals",
        type=int,
        nargs="*",
        help="A year, a number of months, or a negative month offset.",
    )
    args = parser.parse_args()

    now = datetime.datetime.now()
    months_to_render = []

    if len(args.vals) == 0:
        # Single month (Current)
        months_to_render.append(format_month_list(now.year, now.month, now))
        print_months_side_by_side(months_to_render)

    elif len(args.vals) == 1:
        val = args.vals[0]
        if 2 <= val <= 12:
            # Show X number of months starting from current
            curr_y, curr_m = now.year, now.month
            for _ in range(val):
                months_to_render.append(format_month_list(curr_y, curr_m, now))
                curr_m += 1
                if curr_m > 12:
                    curr_m = 1
                    curr_y += 1
            print_months_side_by_side(months_to_render)

        elif -12 < val < 0:
            # render 12 months, starting N months back
            months_back = abs(val)
            total_months = now.year * 12 + now.month - 1
            start_total_months = total_months - months_back

            curr_y = start_total_months // 12
            curr_m = start_total_months % 12 + 1

            for _ in range(12):
                months_to_render.append(format_month_list(curr_y, curr_m, now))
                curr_m += 1
                if curr_m > 12:
                    curr_m = 1
                    curr_y += 1
            print_months_side_by_side(months_to_render)

        elif val >= 1000:
            # Show specific year
            for m in range(1, 13):
                months_to_render.append(format_month_list(val, m, now))
            print_months_side_by_side(months_to_render)

        else:
            print(
                "Error: Input must be a year (e.g. 2025), a count of months (2-12), or a negative month offset (-1 to -11)."
            )

    elif len(args.vals) == 2:
        print("Error: Two arguments are not yet supported.")

    else:
        print("Error: Too many arguments.")


if __name__ == "__main__":
    main()