create_update_kegg_ko.py 3 KB
Newer Older
1
2
3
4
5
6
7
import logging

import requests
from django.core.management.base import BaseCommand
from django.db import IntegrityError

from metagenedb.apps.catalog.models import KeggOrthology
8
from metagenedb.common.utils.chunks import list_chunks
9
10
11
12
13
14
15
16
17
18
19
from metagenedb.common.utils.parsers import KEGGLineParser

logging.basicConfig(format='[%(asctime)s] %(levelname)s:%(name)s:%(message)s')
logger = logging.getLogger(__name__)


class ImportKEGGKO(object):
    KEGG_KO_LIST_API = "http://rest.kegg.jp/list/ko"

    def __init__(self, kegg_ko_list_api=KEGG_KO_LIST_API):
        self.kegg_ko_list_api = kegg_ko_list_api
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
20
21
22
23
24
        self.processed_count = 0
        self.created_count = 0
        self.updated_count = 0
        self.skipped_count = 0
        self.skipped_ids = []
25
26
27
28
29

    def load_all_kegg_ko(self, chunk_size=1000, test=False):
        all_ko_response = requests.get(self.kegg_ko_list_api)
        all_ko_response.raise_for_status()
        all_ko = all_ko_response.text.splitlines()
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
30
        self.total_nb = len(all_ko)
31
        logger.info("%s KEGG KO retrieved from KEGG API...", self.total_nb)
32
        for chunk in list_chunks(all_ko, chunk_size):
33
34
35
36
37
38
            ko_chunk = [KEGGLineParser.ko_list(i) for i in chunk]
            for i in ko_chunk:
                payload = {k: v for k, v in i.items() if v != ""}
                try:
                    kegg = KeggOrthology(**payload)
                    kegg.save()
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
39
                    self.created_count += 1
40
41
42
43
44
45
                except IntegrityError:
                    try:
                        kegg = KeggOrthology.objects.get(function_id=payload.get('function_id'))
                        for k, v in payload.items():
                            setattr(kegg, k, v)
                        kegg.save()
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
46
                        self.updated_count += 1
47
                    except IntegrityError:
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
48
49
50
51
                        self.skipped_ids.append(payload.get('function_id'))
                        self.skipped_count += 1
            self.processed_count += len(ko_chunk)
            logger.info("%s/%s KEGG KO processed so far...", self.processed_count, self.total_nb)
52
53
            if test:
                break
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
54
55
56
57
        logger.info("[DONE] %s/%s KEGG KO created.", self.created_count, self.total_nb)
        logger.info("[DONE] %s/%s KEGG KO updated.", self.updated_count, self.total_nb)
        logger.info("[DONE] %s/%s KEGG KO skipped. List: %s", self.skipped_count, self.total_nb,
                    self.skipped_ids)
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75


class Command(BaseCommand):
    help = 'Create or update all KEGG KO from KEGG API.'

    def add_arguments(self, parser):
        parser.add_argument('--test', action='store_true', help='Run only on first 1000 entries.')

    def set_logger_level(self, verbosity):
        if verbosity > 2:
            logger.setLevel(logging.DEBUG)
        elif verbosity > 1:
            logger.setLevel(logging.INFO)

    def handle(self, *args, **options):
        self.set_logger_level(int(options['verbosity']))
        import_kegg = ImportKEGGKO()
        import_kegg.load_all_kegg_ko(test=options['test'])