scripts/parse_biorxiv_xml.py - master - IvanMysin/Topics

scripts/parse_biorxiv_xml.py
272 строки8 KB
ivan
Working on downloading scripts
17 фев 2026, 15:19
17 фев 2026, 15:192db3ea5
#!/usr/bin/env python
"""
Скрипт для загрузки коллекции документов из bioRxiv XML файлов.

Требования:
- lxml
- sqlite3
"""

import argparse
import os
import sqlite3
from pathlib import Path
from typing import List, Dict, Optional
import sys
from datetime import datetime
from lxml import etree

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.append(str(PROJECT_ROOT))

from src.database_manager import create_articles_table, document_exists, insert_article


def find_xml_files(directory: Path) -> List[Path]:
    """
    Находит все XML файлы в директории.

    Args:
        directory: Путь к директории для поиска XML файлов

    Returns:
        Список путей к XML файлам
    """
    xml_files = []

    print(f"🔍 Поиск XML файлов в директории: {directory}")

    for file in os.listdir(directory):
        if not file.endswith('.xml') or '(' in file or file.startswith('.'):
            continue
        xml_files.append(directory / file)

    print(f"✅ Найдено {len(xml_files)} XML файлов")
    return xml_files


def parse_xml_file(xml_path: Path) -> Dict[str, Optional[str]]:
    """
    Парсит XML файл и извлекает данные статьи.

    Args:
        xml_path: Путь к XML файлу

    Returns:
        Словарь с данными статьи
    """
    tree = etree.parse(str(xml_path))
    root = tree.getroot()

    article_data = {
        'doi': None,
        'title': None,
        'abstract': None,
        'date': None,
        'full_text': None,
        'journal': 'bioRxiv',
        'authors': None,
        'volume': None,
        'issue': None,
        'pages': None,
        'pmid': None,
        'filepath': str(xml_path)
    }

    doi_el = root.xpath('//article-id[@pub-id-type="doi"]')
    if doi_el:
        article_data['doi'] = doi_el[0].text

    title_el = root.xpath('//article-title')
    if title_el:
        article_data['title'] = title_el[0].text

    abstract_el = root.xpath('//abstract //p')
    if abstract_el:
        abstract = etree.tostring(abstract_el[0], pretty_print=True, method='text', encoding='unicode')
        article_data['abstract'] = abstract

    date_el = root.xpath('//date[@date-type="accepted"]')
    if date_el:
        year_el = date_el[0].xpath('//year')
        month_el = date_el[0].xpath('//month')
        day_el = date_el[0].xpath('//day')

        year = year_el[0].text if year_el else ''
        month = month_el[0].text if month_el else ''
        day = day_el[0].text if day_el else ''

        if len(day) == 1:
            day = '0' + day
        if len(month) == 1:
            month = '0' + month

        article_data['date'] = f'{year}-{month}-{day}'.strip('-')

    article_data['full_text'] = parse_full_text(root)

    return article_data


def parse_full_text(root) -> str:
    """
    Парсит полный текст из XML.

    Args:
        root: Корневой элемент XML

    Returns:
        Полный текст статьи
    """
    body_content = root.find('body')

    if body_content is None:
        return ''

    result = []
    for element in body_content.iter():
        if element.tag == 'title':
            current_title = element.text
            if current_title:
                result.append(f"\n{current_title}\n")

        elif element.tag == 'p':
            text_p = etree.tostring(element, pretty_print=True, method='text', encoding='unicode')
            if text_p:
                result.append(f"{text_p.strip()}\n")

    return ' '.join(result)


def save_documents(documents: List[Dict[str, Optional[str]]], db_path: str) -> Dict[str, int]:
    """
    Сохраняет документы в базу данных.

    Args:
        documents: Список документов для сохранения
        db_path: Путь к файлу базы данных

    Returns:
        Статистика по сохраненным и пропущенным документам
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    stats = {
        'processed': 0,
        'saved': 0,
        'skipped': 0,
        'errors': 0
    }

    for doc in documents:
        try:
            stats['processed'] += 1

            doi = doc.get('doi')
            title = doc.get('title')

            if not title or not doc.get('abstract'):
                stats['errors'] += 1
                print(f"❌ Пропуск: нет названия или абстракта")
                continue

            if document_exists(cursor, doi, title):
                stats['skipped'] += 1
                print(f"⏭️ Пропуск дубликата: {title[:50]}...")
                continue

            insert_article(cursor, doc)
            conn.commit()

            stats['saved'] += 1
            print(f"✅ Сохранено: {title[:50]}...")

        except Exception as e:
            stats['errors'] += 1
            print(f"❌ Ошибка при обработке документа: {e}")

    conn.close()
    return stats


def process_directory(xml_dir: Path, db_path: Path) -> Dict[str, int]:
    """
    Обрабатывает директорию с XML файлами.

    Args:
        xml_dir: Директория с XML файлами
        db_path: Путь к базе данных

    Returns:
        Общая статистика обработки
    """
    print("\n" + "="*60)
    print("📁 ЗАГРУЗКА BIORXIV XML ДОКУМЕНТОВ")
    print("="*60)

    create_articles_table(str(db_path))
    print(f"✅ База данных создана/проверена: {db_path}")

    xml_files = find_xml_files(xml_dir)

    if not xml_files:
        print("⚠️ XML файлы не найдены")
        return {'processed': 0, 'saved': 0, 'skipped': 0, 'errors': 0}

    documents = []
    for xml_path in xml_files:
        print(f"Обработка: {xml_path.name}")
        doc_data = parse_xml_file(xml_path)
        documents.append(doc_data)
        print(f'  Символов в абстракте: {len(doc_data.get("abstract", ""))}')
        print(f'  Символов в полном тексте: {len(doc_data.get("full_text", ""))}')
        print('######################')

    stats = save_documents(documents, str(db_path))
    return stats


def main():
    parser = argparse.ArgumentParser(description='Загрузка bioRxiv XML документов в базу данных')
    parser.add_argument('--xml-dir', type=str, required=True,
                        help='Путь к директории с XML файлами')
    parser.add_argument('--db-name', type=str, required=True,
                        help='Название базы данных')

    args = parser.parse_args()

    db_path = Path(args.db_name)
    if not db_path.suffix:
        db_path = db_path.with_suffix('.db')

    if not db_path.is_absolute():
        db_path = PROJECT_ROOT / "data" / db_path

    xml_dir = Path(args.xml_dir)
    if not xml_dir.exists():
        raise FileNotFoundError(f"Директория не найдена: {args.xml_dir}")

    if not xml_dir.is_dir():
        raise NotADirectoryError(f"Указанный путь не является директорией: {args.xml_dir}")

    if not db_path.parent.exists():
        db_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"\n📂 Директория с XML: {xml_dir}")
    print(f"💾 База данных: {db_path}")

    stats = process_directory(xml_dir, db_path)

    print("\n" + "="*60)
    print("📊 ОБЩАЯ СТАТИСТИКА")
    print("="*60)
    print(f"📄 Найдено файлов: {stats['processed']}")
    print(f"✅ Сохранено документов: {stats['saved']}")
    print(f"⏭️ Пропущено (дубликаты): {stats['skipped']}")
    print(f"❌ Ошибок обработки: {stats['errors']}")
    print("="*60)


if __name__ == "__main__":
    main()