cli.py

import argparse
import sys
from pathlib import Path

from extractor.python_extractor import PythonExtractor
from extractor.c_extractor import CExtractor
from extractor.cpp_extractor import CppExtractor
from extractor.java_extractor import JavaExtractor
from extractor.javascript_extractor import JavaScriptExtractor
from extractor.typescript_extractor import TypeScriptExtractor

from utils.extractor_utils import collect_docstrings_in_project, collect_all_docstrings_in_project, save_docstrings_to_json
from utils.git_utils import clone_git_repo, cleanup_git_repo
from utils.logger import logger

# Dictionary mapping language names to extractors
EXTRACTORS = {
    'python': PythonExtractor(),
    'c': CExtractor(),
    'cpp': CppExtractor(),
    'java': JavaExtractor(),
    'javascript': JavaScriptExtractor(),
    'typescript': TypeScriptExtractor(),
}

def main():
    parser = argparse.ArgumentParser(description="Extract docstrings from a project directory or git repository.")
    parser.add_argument(
        'input_dir',
        type=str,
        nargs='?',
        help="Path to the input directory containing source files. If not provided, use --repo-url to clone a repository."
    )
    parser.add_argument(
        'output_file',
        type=str,
        nargs='?',
        default='docstrings.json',
        help="Path to the output JSON file where docstrings will be saved. Defaults to 'docstrings.json'."
    )
    parser.add_argument(
        '--repo-url',
        type=str,
        help="URL of the git repository to clone and extract docstrings from."
    )
    parser.add_argument(
        '--repo-dir',
        type=str,
        default='./repo',
        help="Target directory for cloning the repository. Defaults to './repo'."
    )
    parser.add_argument(
        '--language', '-l',
        type=str,
        choices=list(EXTRACTORS.keys()),
        help="Programming language to extract docstrings for. If not specified, extracts from all supported languages."
    )
    parser.add_argument(
        '--keep-repo',
        action='store_true',
        help="Keep the cloned repository after extraction. By default, the repository is removed."
    )

    args = parser.parse_args()

    output_path = Path(args.output_file)
    input_path = None
    repo_cloned = False

    # Handle git repository cloning
    if args.repo_url:
        if not args.input_dir:
            # Clone repository to specified or default directory
            logger.info(f"Cloning repository from '{args.repo_url}' to '{args.repo_dir}'...")
            success, message = clone_git_repo(args.repo_url, args.repo_dir)

            if not success:
                logger.error(f"Failed to clone repository: {message}")
                sys.exit(1)

            input_path = Path(args.repo_dir)
            repo_cloned = True
            logger.info(f"Repository cloned successfully to '{input_path}'")
        else:
            logger.error("Cannot specify both 'input_dir' and '--repo-url'. Please use only one.")
            sys.exit(1)
    else:
        # Use provided input directory
        if not args.input_dir:
            logger.error("Must specify either 'input_dir' or '--repo-url'")
            sys.exit(1)
        input_path = Path(args.input_dir)

    # Validate input directory exists and is a directory
    if not input_path.exists() or not input_path.is_dir():
        logger.error(f"Input directory '{input_path}' does not exist or is not a directory.")
        sys.exit(1)

    try:
        if args.language:
            # Use specific extractor
            extractor = EXTRACTORS.get(args.language.lower())
            if not extractor:
                logger.error(f"Unsupported language: {args.language}")
                sys.exit(1)
            logger.info(f"Extracting docstrings from {input_path} using {args.language} extractor...")
            results = collect_docstrings_in_project(input_path, extractor)
        else:
            # Use all extractors
            all_extractors = list(EXTRACTORS.values())
            logger.info(f"Extracting docstrings from {input_path} using all supported extractors...")
            results = collect_all_docstrings_in_project(input_path, all_extractors)

        save_docstrings_to_json(results, output_path)
        logger.info(f"Successfully extracted {len(results)} docstrings and saved to {output_path}")

        # Clean up cloned repository if it was cloned and --keep-repo is not specified
        if repo_cloned and not args.keep_repo:
            logger.info(f"Cleaning up cloned repository at '{input_path}'...")
            cleanup_success, cleanup_message = cleanup_git_repo(str(input_path))

            if cleanup_success:
                logger.info("Repository cleaned up successfully")
            else:
                logger.warning(f"Failed to cleanup repository: {cleanup_message}")

    except Exception as e:
        logger.error(f"An error occurred: {e}")
        # Clean up cloned repository on error if it was cloned
        if repo_cloned and not args.keep_repo and input_path and input_path.exists():
            logger.info("Cleaning up cloned repository due to error...")
            cleanup_git_repo(str(input_path))
        sys.exit(1)

if __name__ == "__main__":
    main()