create_update_eggnog.py 5.07 KB
Newer Older
1
2
3
import logging

from django.core.management.base import BaseCommand
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
4
from django.core.exceptions import ValidationError
5

6
from metagenedb.apps.catalog.models import EggNOG, EggNogFunctionalCategory
7
from metagenedb.common.utils.chunks import file_len
8
from metagenedb.common.utils.parsers import EggNOGAnnotationLineParser
9
10
11
12
13

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


14
class ImportEggNOG(object):
15
16
17

    def __init__(self, file_path):
        self.annotation_file = file_path
18
        self.eggnog_parser = EggNOGAnnotationLineParser()
19
20
21
22
23
        self.processed_count = 0
        self.created_count = 0
        self.updated_count = 0
        self.skipped_count = 0
        self.skipped_ids = []
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
24
        self.skipped_errors = []
25
26

    def _build_functional_category_dict(self):
27
        all_categories = EggNogFunctionalCategory.objects.all()
28
29
        if not all_categories:
            raise Exception("You need to create Functional categories first.")
30
        self.functional_cat_instances = {cat.category_id: cat for cat in all_categories}
31

32
33
    def add_functional_categories(self, eggnog, functional_category_keys):
        for key in functional_category_keys:
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
34
35
36
37
            if key not in self.functional_cat_instances.keys():
                logger.warning("%s is not a valid functional categories. 'S' used instead for %s",
                               key, eggnog.function_id)
                key = 'S'
38
            eggnog.functional_categories.add(self.functional_cat_instances[key])
39

40
    def load_all(self, test=False, version=None):
41
42
        self._build_functional_category_dict()
        self.total_eggnog_nb = file_len(self.annotation_file)
43
        logger.info("Preparing to load %s items...", self.total_eggnog_nb)
44
45
46
        with open(self.annotation_file, "r") as file:
            for line in file:
                eggnog_dict = self.eggnog_parser.get_dict(line)
47
                functional_category_keys = eggnog_dict.pop('functional_categories')  # link later
48
                payload = {k: v for k, v in eggnog_dict.items() if v != ""}
49
                payload['version'] = version
50
                eggnog = None
51
                try:
52
                    eggnog = EggNOG(**payload)
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
53
                    eggnog.full_clean()
54
55
                    eggnog.save()
                    self.created_count += 1
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
56
57
58
                except ValidationError as validation_error:
                    if 'function_id' in validation_error.error_dict.keys():
                        try:
59
                            eggnog = EggNOG.objects.get(function_id=payload.get('function_id'))
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
60
61
62
63
64
65
66
67
68
69
                            for k, v in payload.items():
                                setattr(eggnog, k, v)
                            eggnog.full_clean()
                            eggnog.save()
                            self.updated_count += 1
                        except ValidationError as validation_error:
                            logger.error(validation_error)
                            self.skipped_errors.append(validation_error)
                            self.skipped_ids.append(payload.get('function_id'))
                            self.skipped_count += 1
70
                            eggnog = None
Kenzo-Hugo Hillion's avatar
Kenzo-Hugo Hillion committed
71
72
73
                    else:
                        logger.error(validation_error)
                        self.skipped_errors.append(validation_error)
74
75
                        self.skipped_ids.append(payload.get('function_id'))
                        self.skipped_count += 1
76
                        eggnog = None
77
78
79
                if eggnog is not None:
                    self.add_functional_categories(eggnog, functional_category_keys)
                    eggnog.save()
80
81
                self.processed_count += 1
                if self.processed_count % 1000 == 0:
82
                    logger.info("%s/%s EggNOG processed so far...", self.processed_count, self.total_eggnog_nb)
83
84
                    if test:
                        break
85
86
87
        logger.info("[DONE] %s/%s EggNOG created.", self.created_count, self.total_eggnog_nb)
        logger.info("[DONE] %s/%s EggNOG updated.", self.updated_count, self.total_eggnog_nb)
        logger.info("[DONE] %s/%s EggNOG skipped. List: %s", self.skipped_count, self.total_eggnog_nb,
88
89
90
91
                    self.skipped_ids)


class Command(BaseCommand):
92
    help = 'Create or update all EggNOG entries from annotations.tsv file.'
93
94

    def add_arguments(self, parser):
95
        parser.add_argument('annotation', help='annotations.tsv file from EggNOG')
96
        parser.add_argument('--test', action='store_true', help='Run only on first 1000 entries.')
97
        parser.add_argument('--eggnog_version', choices=['3.0', '5.0'], required=True, help='version of EggNOG.')
98
99
100
101
102
103
104
105
106

    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']))
107
        import_eggnog = ImportEggNOG(options['annotation'])
108
        import_eggnog.load_all(test=options['test'], version=options['eggnog_version'])