storage.py

"""Simple JSON-based storage for tracking PRs."""
import json
import os
from datetime import datetime
from typing import Dict, List, Optional
from config import PR_TRACKING_FILE, MAX_TRACKED_PRS


class PRStorage:
    """Lightweight storage for tracking recent PRs."""
    
    def __init__(self, storage_file: str = PR_TRACKING_FILE):
        self.storage_file = storage_file
        self._ensure_file_exists()
    
    def _ensure_file_exists(self):
        """Create storage file if it doesn't exist."""
        if not os.path.exists(self.storage_file):
            with open(self.storage_file, 'w') as f:
                json.dump({"prs": []}, f)
    
    def _read_data(self) -> Dict:
        """Read data from storage file."""
        try:
            with open(self.storage_file, 'r') as f:
                return json.load(f)
        except (json.JSONDecodeError, FileNotFoundError):
            return {"prs": []}
    
    def _write_data(self, data: Dict):
        """Write data to storage file."""
        with open(self.storage_file, 'w') as f:
            json.dump(data, f, indent=2)
    
    def add_pr(self, pr_data: Dict):
        """Add a PR to tracking storage."""
        data = self._read_data()
        
        # Add timestamp
        pr_data["tracked_at"] = datetime.utcnow().isoformat()
        
        # Add to beginning of list
        data["prs"].insert(0, pr_data)
        
        # Keep only recent PRs
        data["prs"] = data["prs"][:MAX_TRACKED_PRS]
        
        self._write_data(data)
    
    def get_pr(self, repo_full_name: str, pr_number: int) -> Optional[Dict]:
        """Get a specific PR from storage."""
        data = self._read_data()
        for pr in data["prs"]:
            if pr.get("repo") == repo_full_name and pr.get("pr_number") == pr_number:
                return pr
        return None
    
    def get_all_prs(self) -> List[Dict]:
        """Get all tracked PRs."""
        data = self._read_data()
        return data.get("prs", [])
    
    def update_pr(self, repo_full_name: str, pr_number: int, update_data: Dict):
        """Update a PR in storage."""
        data = self._read_data()
        for i, pr in enumerate(data["prs"]):
            if pr.get("repo") == repo_full_name and pr.get("pr_number") == pr_number:
                data["prs"][i].update(update_data)
                data["prs"][i]["updated_at"] = datetime.utcnow().isoformat()
                self._write_data(data)
                return True
        return False