che-astronvim-config/lua/plugins/pendulum/branch_time_calculator.py

import csv
import datetime
import os
import subprocess
import argparse
from typing import Optional, Dict, List


class BranchTimeCalculator:
    """
    A class to calculate time spent on a Git branch with exact project matching.
    Handles lowercase 'true' for active status and flexible CSV parsing.
    """

    def __init__(self, csv_file: Optional[str] = None, target_branch: Optional[str] = None, debug: bool = False):
        """
        Initialize the calculator with optional CSV file, target branch and debug mode.
        
        Args:
            csv_file (Optional[str]): Path to CSV log file. Defaults to ~/pendulum-log.csv.
            target_branch (Optional[str]): Name of the target branch. Defaults to current branch.
            debug (bool): If True, prints debug information. Defaults to False.
        """
        self.csv_file = csv_file or os.path.expanduser("~/pendulum-log.csv")
        self.target_branch = target_branch
        self.debug = debug
        self.current_project = self._get_project_name()

    def _get_project_name(self) -> str:
        """
        Get the project name from Git config or fall back to directory name.
        
        Returns:
            str: The project name.
        """
        try:
            return subprocess.check_output(
                ["git", "config", "--get", "remote.origin.url"],
                stderr=subprocess.DEVNULL,
                text=True
            ).strip().split('/')[-1].replace('.git', '')
        except subprocess.CalledProcessError:
            return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def _get_current_branch(self) -> Optional[str]:
        """
        Get the current Git branch name.
        
        Returns:
            Optional[str]: The branch name or None if not in a Git repo.
        """
        try:
            return subprocess.check_output(
                ["git", "branch", "--show-current"],
                stderr=subprocess.DEVNULL,
                text=True
            ).strip()
        except subprocess.CalledProcessError:
            print("Error: Not a Git repository or branch not found.")
            return None

    def _parse_log_entry(self, row: List[str]) -> Dict[str, str]:
        """
        Parse a CSV row into a log entry dictionary.
        
        Args:
            row (List[str]): A row from the CSV log file.
            
        Returns:
            Dict[str, str]: A dictionary representing the log entry.
        """
        return {
            'active': row[0].lower().strip(),
            'branch': row[1].strip(),
            'directory': row[2].strip(),
            'file': row[3].strip(),
            'filetype': row[4].strip(),
            'project': row[5].strip(),
            'time': row[6].strip()
        }

    def _calculate_time_deltas(self, log_entries: List[Dict[str, str]]) -> Dict[str, datetime.timedelta]:
        """
        Calculate total and active time from log entries.
        
        Args:
            log_entries (List[Dict[str, str]]): List of parsed log entries.
            
        Returns:
            Dict[str, datetime.timedelta]: A dictionary with 'total_time' and 'active_time'.
        """
        total_time = datetime.timedelta()
        active_time = datetime.timedelta()
        prev_time = None

        for entry in log_entries:
            try:
                current_time = datetime.datetime.strptime(entry['time'], '%Y-%m-%d %H:%M:%S')
            except ValueError as e:
                print(f"[ERROR] Time parse failed: {entry['time']}. Error: {e}")
                continue

            if prev_time is not None:
                time_diff = current_time - prev_time
                total_time += time_diff
                if entry['active'] == 'true':
                    active_time += time_diff

            prev_time = current_time

        return {'total_time': total_time, 'active_time': active_time}

    def calculate(self) -> Optional[Dict[str, any]]:
        """
        Calculate time spent on the target branch.
        
        Returns:
            Optional[Dict[str, any]]: A dictionary with branch, project, total_time, active_time, and active_percentage.
                                     Returns None if calculation fails.
        """
        if self.target_branch is None:
            self.target_branch = self._get_current_branch()
            if self.target_branch is None:
                return None

        log_entries = []
        try:
            with open(self.csv_file, mode='r') as file:
                for line in file:
                    row = list(csv.reader([line.strip()], delimiter=',', quotechar='"'))[0]
                    if len(row) != 7:
                        if self.debug:
                            print(f"[DEBUG] Skipping malformed line: {line}")
                        continue

                    log_entry = self._parse_log_entry(row)
                    if log_entry['branch'] == self.target_branch and log_entry['project'] == self.current_project:
                        log_entries.append(log_entry)

        except FileNotFoundError:
            print(f"Error: Log file not found at {self.csv_file}")
            return None

        time_deltas = self._calculate_time_deltas(log_entries)
        total_seconds = time_deltas['total_time'].total_seconds()
        active_percentage = (time_deltas['active_time'].total_seconds() / total_seconds) * 100 if total_seconds > 0 else 0

        if self.debug:
            print(f"\nTime spent on branch '{self.target_branch}' in project '{self.current_project}':")
            print(f"- Total time: {time_deltas['total_time']}")
            print(f"- Active time: {time_deltas['active_time']}")
            print(f"- Active percentage: {active_percentage:.2f}%")
        else:
            print(f"In project '{self.current_project}' on branch '{self.target_branch}' you spent actively working: {time_deltas['active_time']}")

        return {
            'branch': self.target_branch,
            'project': self.current_project,
            'total_time': time_deltas['total_time'],
            'active_time': time_deltas['active_time'],
            'active_percentage': active_percentage
        }


def main():
    """Main function to parse arguments and run the calculator."""
    parser = argparse.ArgumentParser(description='Calculate branch activity time.')
    parser.add_argument('--csv-file', type=str, help='Path to CSV log (default: ~/pendulum-log.csv)')
    parser.add_argument('--branch', type=str, help='Branch name (default: current)')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')
    args = parser.parse_args()

    calculator = BranchTimeCalculator(args.csv_file, args.branch, args.debug)
    result = calculator.calculate()

    if not result:
        print("\nCalculation failed - check debug output above.")


if __name__ == '__main__':
    main()