collect_data_pygithub.py

from github import Github
import requests
import logging
from datetime import datetime
from pathlib import Path
import hashlib
import time
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import random
import os
import shutil
import tempfile

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 配置 GitHub API 访问
def get_github_token():
    """从环境变量获取GitHub token"""
    token = os.getenv('GITHUB_TOKEN')
    if not token:
        token_file = Path.home() / '.github_token'
        if token_file.exists():
            return token_file.read_text().strip()
    return token

github_token = get_github_token()
if not github_token:
    logging.error('GitHub token not found. Please set GITHUB_TOKEN environment variable or create ~/.github_token file')
    exit(1)

# 配置重试策略
retry_strategy = Retry(
    total=10,  # 增加总重试次数
    backoff_factor=2,  # 增加退避因子
    status_forcelist=[500, 502, 503, 504, 429],  # 添加429（速率限制）状态码
    allowed_methods=["GET", "POST"]  # 允许重试的HTTP方法
)

# 配置 GitHub 客户端
try:
    github = Github(
        github_token,
        per_page=50,  # 减少每页请求数量
        timeout=60,  # 增加超时时间
        retry=retry_strategy,
        verify=True  # 确保SSL验证
    )
    logging.info('Successfully authenticated with GitHub')
except Exception as e:
    logging.error(f'Error authenticating with GitHub: {str(e)}')
    exit(1)

# 项目列表
projects = [
    {'name': 'THUDM/CodeGeeX', 'size': None}  # size 将在运行时确定
]

sizemap = {
    'S': (0, 9999),
    'M': (10000, 99999),
    'L': (100000, 999999)
}

def count_python_lines(repo):
    """统计仓库中Python文件的总行数"""
    total_lines = 0
    try:
        # 获取仓库的默认分支
        default_branch = repo.default_branch
        # 获取仓库内容
        contents = repo.get_contents("", ref=default_branch)
        
        def process_contents(contents):
            nonlocal total_lines
            for content in contents:
                if content.type == "dir":
                    # 如果是目录，递归处理
                    sub_contents = repo.get_contents(content.path, ref=default_branch)
                    process_contents(sub_contents)
                elif content.name.endswith('.py'):
                    # 如果是Python文件，获取内容并计算行数
                    try:
                        file_content = content.decoded_content.decode('utf-8')
                        total_lines += len(file_content.splitlines())
                    except Exception as e:
                        logging.warning(f"Error processing file {content.path}: {str(e)}")
        
        process_contents(contents)
        return total_lines
    except Exception as e:
        logging.error(f"Error counting lines: {str(e)}")
        return 0

def get_project_size(repo):
    """根据Python文件行数确定项目大小"""
    total_lines = count_python_lines(repo)
    for size, (min_lines, max_lines) in sizemap.items():
        if min_lines <= total_lines < max_lines:
            return size
    return 'Unknown'

# 存储收集到的数据
data = []

def random_sleep():
    """随机延迟，避免固定间隔"""
    time.sleep(random.uniform(1.0, 2.0))

def generate_task_id(commit_sha, repo_name):
    """生成唯一的task_id"""
    return hashlib.md5(f"{repo_name}_{commit_sha}".encode()).hexdigest()[:8]

def download_repo_locally(repo_name, branch='main'):
    """下载仓库到本地临时目录"""
    temp_dir = Path(tempfile.mkdtemp())
    repo_dir = temp_dir / repo_name.split('/')[-1]
    
    try:
        # 使用git clone命令下载仓库
        clone_cmd = f"git clone --depth 1 --branch {branch} https://github.com/{repo_name}.git {repo_dir}"
        os.system(clone_cmd)
        return repo_dir
    except Exception as e:
        logging.error(f"Error cloning repository {repo_name}: {str(e)}")
        return None

def get_commit_details_local(repo_dir, commit_sha):
    """从本地仓库获取提交详情"""
    try:
        # 切换到指定目录
        os.chdir(repo_dir)
        
        # 获取提交信息
        commit_info = os.popen(f"git show {commit_sha} --name-only --pretty=format:'%s'").read()
        
        # 获取修改的文件
        changed_files = os.popen(f"git show --name-only --pretty=format: {commit_sha}").read().strip().split('\n')
        
        # 获取diff
        diff = os.popen(f"git show {commit_sha}").read()
        
        return {
            'commit_sha': commit_sha,
            'commit_message': commit_info.split('\n')[0],
            'affected_files': [f for f in changed_files if f.endswith('.py')],
            'diff': diff
        }
    except Exception as e:
        logging.error(f"Error getting commit details locally: {str(e)}")
        return None

def determine_change_type(commit_message, affected_files):
    """判断变更类型
    返回: 'bugfix', 'refactoring', 'api', 或 'feature'
    """
    commit_message = commit_message.lower()
    
    # 检查是否是bug修复
    bug_keywords = ['fix', 'bug', 'issue', 'error', 'exception', 'fail', 'crash', 'defect']
    if any(keyword in commit_message for keyword in bug_keywords):
        return 'bugfix'
    
    # 检查是否是重构
    refactor_keywords = ['refactor', 'restructure', 'reorganize', 'cleanup', 'optimize']
    if any(keyword in commit_message for keyword in refactor_keywords):
        return 'refactoring'
    
    # 检查是否是API变更
    api_keywords = ['api', 'interface', 'contract', 'signature', 'method', 'function']
    if any(keyword in commit_message for keyword in api_keywords):
        return 'api'
    
    # 检查是否是功能完善
    feature_keywords = ['add', 'implement', 'feature', 'support', 'enhance', 'improve']
    if any(keyword in commit_message for keyword in feature_keywords):
        return 'feature'
    
    return 'other'

def collect_data(repo):
    try:
        logging.info(f'Collecting data for repository: {repo.full_name}')
        
        # 下载仓库到本地
        repo_dir = download_repo_locally(repo.full_name)
        if not repo_dir:
            logging.error(f"Failed to download repository {repo.full_name}")
            return
            
        # 确定项目大小
        project_size = get_project_size(repo)
        logging.info(f'Project size determined as: {project_size}')
        
        # 获取提交历史
        commits = os.popen(f"cd {repo_dir} && git log --pretty=format:'%H'").read().strip().split('\n')
        
        for commit_sha in commits:
            try:
                # 从本地获取提交详情
                commit_details = get_commit_details_local(repo_dir, commit_sha)
                if not commit_details:
                    continue
                    
                # 如果这个提交没有修改Python文件，跳过
                if not commit_details['affected_files']:
                    continue
                    
                task_id = generate_task_id(commit_sha, repo.full_name)
                
                data.append({
                    'task_id': task_id,
                    'project': {
                        'repo_url': repo.html_url,
                        'size_category': project_size,
                        'language': 'Python',
                        'total_python_lines': count_python_lines(repo)
                    },
                    'change': {
                        'description': commit_details['commit_message'],
                        'change_type': determine_change_type(commit_details['commit_message'], commit_details['affected_files']),
                        'branch': repo.default_branch,
                        'affected_files': commit_details['affected_files'],
                        'diff': commit_details['diff'],
                        'commit': {
                            'before': None,
                            'after': commit_sha
                        },
                        'timestamp': datetime.now().isoformat()
                    }
                })
                
                logging.info(f'Processed commit: {commit_sha}')
                random_sleep()
                
            except Exception as e:
                logging.error(f'Error processing commit {commit_sha}: {str(e)}')
                continue
                
    except Exception as e:
        logging.error(f'Error collecting data from repository {repo.full_name}: {str(e)}')
    finally:
        # 清理临时目录
        if repo_dir and os.path.exists(repo_dir):
            shutil.rmtree(repo_dir.parent)

for project in projects:
    try:
        logging.info(f'Processing project: {project["name"]}')
        repo = github.get_repo(project['name'])
        if not repo:
            logging.error(f'Repository not found: {project["name"]}')
            continue
        collect_data(repo)
    except Exception as e:
        logging.error(f'Error processing project {project["name"]}: {str(e)}')

# 保存数据到文件
import json
try:
    output_file = 'collected_data_pygithub.json'
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    logging.info(f'Successfully saved data to {output_file}')
    logging.info(f'Total records collected: {len(data)}')
except Exception as e:
    logging.error(f'Error saving data to file: {str(e)}')
    exit(1)