Add a fiscal year column to a CSV file based on dates.

add_fiscal_year_to_csv.py

#!/usr/bin/env python3
"""
Add a fiscal year column to a CSV file based on dates,
supporting both US and Japanese fiscal calendar systems.
"""
import argparse
from datetime import datetime
from pathlib import Path

import pandas as pd


def calculate_fiscal_year(date_iso: str, system: str = "us") -> int:
    """
    Determines the fiscal year for a given date, based on the specified fiscal system.

    Calculates the fiscal year for a date provided in ISO format according to the
    specified fiscal system. Supports 'us' and 'jp' fiscal systems, where the U.S.
    fiscal year ends on September 30 and the Japanese fiscal year ends on March 31.
    Raises an error for unsupported fiscal systems.

    Args:
        date_iso (str): The date in ISO format (YYYY-MM-DD).
        system (str): The fiscal system to use. Defaults to "us". Supported values
            are:
            - "us": Fiscal year ends September 30.
            - "jp": Fiscal year ends March 31.

    Returns:
        int: The fiscal year corresponding to the provided date and system.

    Raises:
        ValueError: If an unsupported fiscal system is provided.
    """
    date_obj = datetime.fromisoformat(date_iso).date()

    if system == "us":  # Fiscal year ends Sep 30 → starts Oct 1 previous year
        return date_obj.year + 1 if date_obj.month >= 10 else date_obj.year
    elif system == "jp":  # Fiscal year ends Mar 31 → starts Apr 1 previous year
        return date_obj.year if date_obj.month >= 4 else date_obj.year - 1
    else:
        raise ValueError(f"Unsupported fiscal system: {system}")


def add_fiscal_year_to_csv(input_csv: Path, output_csv: Path, date_column: str, fiscal_system: str):
    df = pd.read_csv(input_csv)

    if date_column not in df.columns:
        raise ValueError(f"Column '{date_column}' not found in CSV.")

    df['fiscal_year'] = df[date_column].apply(lambda x: calculate_fiscal_year(x, fiscal_system))
    df.to_csv(output_csv, index=False)
    print(f"Output saved to: {output_csv} (Fiscal System: {fiscal_system.upper()})")


def main():
    parser = argparse.ArgumentParser(description="Add fiscal year column to a CSV.")
    parser.add_argument("input_csv", type=Path, help="Path to input CSV file")
    parser.add_argument("output_csv", type=Path, help="Path to output CSV file")
    parser.add_argument("--date-column", default="date", help="Name of date column (default: 'date')")
    parser.add_argument(
        "--fiscal-system",
        choices=["us", "jp"],
        default="us",
        help="Fiscal system to use: 'us' (Sep-30 end) or 'ja' (Mar-31 end). Default is 'us'."
    )

    args = parser.parse_args()
    add_fiscal_year_to_csv(args.input_csv, args.output_csv, args.date_column, args.fiscal_system)


if __name__ == "__main__":
    main()