kicad-manager/variant_manager.py

#!/usr/bin/env python3
"""
variant_manager.py
==================
Manage KiCad design variants - track which parts are fitted/unfitted in different variants.

Variants are stored in a JSON file alongside the project.
"""

import json
import sys
from pathlib import Path
from typing import Dict, List, Set


class VariantManager:
    """Manages design variants for a KiCad project"""

    def __init__(self, project_path: str):
        """
        Initialize variant manager for a project.

        Args:
            project_path: Path to the .kicad_pro or .kicad_sch file
        """
        self.project_path = Path(project_path)
        self.project_dir = self.project_path.parent
        self.project_name = self.project_path.stem

        # Variants file stored alongside project
        self.variants_file = self.project_dir / f"{self.project_name}.variants.json"
        self.variants = self._load_variants()

    def _load_variants(self) -> Dict:
        """Load variants from JSON file"""
        if self.variants_file.exists():
            with open(self.variants_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        else:
            # Default structure
            return {
                "meta": {
                    "version": 2,  # Version 2 uses UUIDs instead of references
                    "active_variant": "default"
                },
                "variants": {
                    "default": {
                        "name": "default",
                        "description": "Default variant - all parts fitted",
                        "dnp_parts": []  # List of UUIDs that are DNP (Do Not Place)
                    }
                }
            }

    def _save_variants(self):
        """Save variants to JSON file"""
        with open(self.variants_file, 'w', encoding='utf-8') as f:
            json.dump(self.variants, f, indent=2)

    def get_variants(self) -> Dict:
        """Get all variants"""
        return self.variants["variants"]

    def get_active_variant(self) -> str:
        """Get the name of the active variant"""
        return self.variants["meta"]["active_variant"]

    def create_variant(self, name: str, description: str = "", based_on: str = None) -> bool:
        """
        Create a new variant.

        Args:
            name: Variant name
            description: Variant description
            based_on: Name of variant to copy from (None = start empty)

        Returns:
            True if created, False if already exists
        """
        if name in self.variants["variants"]:
            return False

        if based_on and based_on in self.variants["variants"]:
            # Copy DNP list from base variant
            dnp_parts = self.variants["variants"][based_on]["dnp_parts"].copy()
        else:
            dnp_parts = []

        self.variants["variants"][name] = {
            "name": name,
            "description": description,
            "dnp_parts": dnp_parts
        }

        self._save_variants()
        return True

    def delete_variant(self, name: str) -> bool:
        """
        Delete a variant.

        Args:
            name: Variant name

        Returns:
            True if deleted, False if doesn't exist or is active
        """
        if name not in self.variants["variants"]:
            return False

        if name == self.variants["meta"]["active_variant"]:
            return False  # Can't delete active variant

        if name == "default":
            return False  # Can't delete default variant

        del self.variants["variants"][name]
        self._save_variants()
        return True

    def set_active_variant(self, name: str) -> bool:
        """
        Set the active variant.

        Args:
            name: Variant name

        Returns:
            True if set, False if variant doesn't exist
        """
        if name not in self.variants["variants"]:
            return False

        self.variants["meta"]["active_variant"] = name
        self._save_variants()
        return True

    def set_part_dnp(self, variant_name: str, uuid: str, is_dnp: bool) -> bool:
        """
        Set whether a part is DNP (Do Not Place) in a variant.

        Args:
            variant_name: Variant name
            uuid: Component UUID (e.g., "681abb84-6eb2-4c95-9a2f-a9fc19a34beb")
            is_dnp: True to mark as DNP, False to mark as fitted

        Returns:
            True if successful, False if variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return False

        dnp_list = self.variants["variants"][variant_name]["dnp_parts"]

        if is_dnp:
            if uuid not in dnp_list:
                dnp_list.append(uuid)
                dnp_list.sort()  # Keep sorted
        else:
            if uuid in dnp_list:
                dnp_list.remove(uuid)

        self._save_variants()
        return True

    def get_dnp_parts(self, variant_name: str) -> List[str]:
        """
        Get list of DNP parts for a variant.

        Args:
            variant_name: Variant name

        Returns:
            List of UUIDs, or empty list if variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return []

        return self.variants["variants"][variant_name]["dnp_parts"].copy()

    def is_part_dnp(self, variant_name: str, uuid: str) -> bool:
        """
        Check if a part is DNP in a variant.

        Args:
            variant_name: Variant name
            uuid: Component UUID

        Returns:
            True if DNP, False if fitted or variant doesn't exist
        """
        if variant_name not in self.variants["variants"]:
            return False

        return uuid in self.variants["variants"][variant_name]["dnp_parts"]


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python variant_manager.py <project.kicad_pro>")
        sys.exit(1)

    manager = VariantManager(sys.argv[1])

    # Print current variants
    print(f"Project: {manager.project_name}")
    print(f"Active variant: {manager.get_active_variant()}")
    print("\nVariants:")
    for name, variant in manager.get_variants().items():
        dnp_count = len(variant["dnp_parts"])
        print(f"  {name}: {variant['description']} ({dnp_count} DNP parts)")