OpenICLR/data_storage.py at main · wonvicky/OpenICLR · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
"""
数据存储模块
用于将爬取的数据保存为JSON和CSV格式
"""

import json
import pandas as pd
import os
from typing import Dict, List
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DataStorage:
    """数据存储类"""

    def __init__(self, output_dir: str = "data"):
        """
        初始化数据存储

        Args:
            output_dir: 输出目录
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def save_json(self, data: Dict, filename: str):
        """
        保存数据为JSON格式

        Args:
            data: 要保存的数据
            filename: 文件名
        """
        filepath = self.output_dir / filename
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        logger.info(f"JSON数据已保存到: {filepath}")

    def save_csv(self, data: List[Dict], filename: str):
        """
        保存数据为CSV格式

        Args:
            data: 要保存的数据列表
            filename: 文件名
        """
        if not data:
            logger.warning(f"数据为空，跳过CSV保存: {filename}")
            return

        df = pd.DataFrame(data)
        filepath = self.output_dir / filename
        df.to_csv(filepath, index=False, encoding='utf-8-sig')
        logger.info(f"CSV数据已保存到: {filepath}")

    def save_iclr_data(self, data: Dict, year: int):
        """
        保存ICLR数据（包括论文、评审、审稿人）

        Args:
            data: ICLR数据字典
            year: 年份
        """
        # 保存完整JSON
        self.save_json(data, f"iclr_{year}_complete.json")

        # 保存论文数据
        papers = data.get('papers', [])
        if papers:
            # 展开论文数据（移除嵌套的reviews）
            papers_flat = []
            for paper in papers:
                paper_flat = paper.copy()
                if 'reviews' in paper_flat:
                    paper_flat['num_reviews'] = len(paper_flat['reviews'])
                    del paper_flat['reviews']
                papers_flat.append(paper_flat)

            self.save_csv(papers_flat, f"iclr_{year}_papers.csv")
            self.save_json({'papers': papers}, f"iclr_{year}_papers.json")

        # 保存评审数据
        reviews = data.get('reviews', [])
        if reviews:
            self.save_csv(reviews, f"iclr_{year}_reviews.csv")
            self.save_json({'reviews': reviews}, f"iclr_{year}_reviews.json")

        # 保存审稿人数据
        reviewers = data.get('reviewers', [])
        if reviewers:
            self.save_csv(reviewers, f"iclr_{year}_reviewers.csv")
            self.save_json({'reviewers': reviewers}, f"iclr_{year}_reviewers.json")

        # 保存统计信息
        stats = {
            'year': year,
            'total_papers': data.get('total_papers', 0),
            'total_reviews': data.get('total_reviews', 0),
            'total_reviewers': data.get('total_reviewers', 0)
        }
        self.save_json(stats, f"iclr_{year}_stats.json")
        logger.info(f"所有ICLR {year}数据已保存到: {self.output_dir}")