update_mrc_files.py 5.74 KB
Newer Older
1
2
3
4
from django.core.management.base import AppCommand, CommandError
from ippidb.management.commands import TaskOutWrapper
from ippidb.models import Chain, InteractFile
from django.core.files import File
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
5
from celery import states
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
6
import os
7
import sys
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
8
9
10
11
12
13
14
15
16
17
18
19
20
import glob


def dir_path(string):
    """
    Test if a path is a directory
    """
    if os.path.isdir(string):
        return string
    else:
        raise NotADirectoryError(string)


21
22
23
24
25
def isolevel_float(arg):
    """ Type function for argparse - a float within some predefined bounds """
    try:
        f = float(arg)
    except ValueError:
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
26
        raise CommandError("Must be a floating point number")
27
    if f < 0 or f > 1.0:
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
28
        raise CommandError("Argument must be < 1.0 and > 0")
29
30
31
    return f


Fabien  MAREUIL's avatar
Fabien MAREUIL committed
32
33
34
35
36
class Command(AppCommand):
    """
    Command to update mrc files
    """

Fabien  MAREUIL's avatar
Fabien MAREUIL committed
37
    help = "liste all mrc files in a directory and update files in the database"
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
38

39
    def __init__(
40
41
42
43
44
45
        self,
        stdout=None,
        stderr=None,
        no_color=False,
        force_color=False,
        task=None,
46
47
48
49
    ):
        super(Command, self).__init__(
            stdout=stdout, stderr=stderr, no_color=no_color, force_color=force_color
        )
50
51
52
53
        if task:
            task.update_state(state=states.STARTED)
            self.stdout = TaskOutWrapper(stdout or sys.stdout, task=task, std_out=True)
            self.stderr = TaskOutWrapper(stderr or sys.stderr, task=task, std_err=True)
54

Fabien  MAREUIL's avatar
Fabien MAREUIL committed
55
56
    def add_arguments(self, parser):
        parser.add_argument(
57
58
59
60
            "-mp",
            "--mrcpath",
            type=dir_path,
            help="root directory with all mrc files",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
61
62
        )
        parser.add_argument(
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
63
64
            "-pa",
            "--pattern",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
65
66
            dest="pattern",
            help="pattern for select files",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
67
            default="*.mrc",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
68
69
        )
        parser.add_argument(
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
70
71
            "-t",
            "--type",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
72
73
74
            dest="type",
            choices=["druggability", "interactibility"],
            default="druggability",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
75
            help="files type, druggability or interactibility",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
76
77
        )
        parser.add_argument(
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
78
79
            "-l",
            "--label",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
80
            dest="label",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
81
            help="label for files, used as label in the interface",
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
82
        )
83
84
85
86
87
88
89
90
        parser.add_argument(
            "-di",
            "--default_isolevel",
            dest="default_isolevel",
            help="isolevel default value",
            default=0.7,
            type=isolevel_float,
        )
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
91

Fabien  MAREUIL's avatar
Fabien MAREUIL committed
92
93
94
95
96
97
98
99
    def handle(self, *args, **options):
        """
        Perform the command's actions
        """
        typefile = options["type"]
        pattern = options["pattern"]
        mrcpath = options["mrcpath"]
        label = options.get("label")
100
101
        if typefile == "interactibility" and not label:
            raise CommandError("Error: Interactibility need a label")
102
103
104
        if typefile == "druggability":
            storage_path = Chain.mrc_file.field.storage.path("content/mrc_files")
        else:
105
106
107
            storage_path = InteractFile.interact_file.field.storage.path(
                "content/interact_files"
            )
Fabien  MAREUIL's avatar
typo if    
Fabien MAREUIL committed
108
        if not os.path.exists(storage_path):
109
            os.makedirs(storage_path)
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
110
        list_files = glob.glob(os.path.join(mrcpath, pattern))
111
112
113
114
115
116
117
118
        for data in list_files:
            self.stdout.write(
                self.style.WARNING(
                    "Processing mrc file {}".format(os.path.basename(data))
                )
            )
            splitname = os.path.basename(data).split(".")[0].split("-")
            if len(splitname) == 5:
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
119
                chain = splitname[-1].split("_")[0]
120
121
                pdb = splitname[0]
            elif len(splitname) == 3:
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
122
                chain = splitname[1].split("_")[0]
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
                pdb = splitname[0]
            else:
                raise CommandError(
                    "Error: Wrong name file {},"
                    "required format is pdb-chains-uniprot_chain1-uniprot_chain2-chain.mrc"
                    "or pdb-chain-uniport.mrc".format(os.path.basename(data))
                )
            dbchains = Chain.objects.filter(pdb_chain_id=chain, pdb__code=pdb)
            if dbchains.count() > 1:
                raise CommandError(
                    "More than one chain with pdb: {} chain: {}".format(pdb, chain)
                )
            elif dbchains.count() == 1:
                if typefile == "druggability":
                    with open(data, "rb") as f:
                        dbchains[0].mrc_file.delete()
                        dbchains[0].mrc_file.save(os.path.basename(data), File(f))
140
141
                        dbchains[0].default_isolevel = options["default_isolevel"]
                        dbchains[0].save()
142
143
                else:
                    interactfile, created = InteractFile.objects.get_or_create(
144
145
                        chain=dbchains[0],
                        label=label,
146
147
148
149
                    )
                    with open(data, "rb") as f:
                        interactfile.interact_file.delete()
                        interactfile.interact_file.save(os.path.basename(data), File(f))
Fabien  MAREUIL's avatar
Fabien MAREUIL committed
150
                    interactfile.default_isolevel = options["default_isolevel"]
151
                    interactfile.save()
152
153
154
155
156
157
158
159
160
161
162
                self.stdout.write(
                    self.style.SUCCESS(
                        "File mrc {}, have been saved in Chain, pdb: {}, chain: {}".format(
                            os.path.basename(data), pdb, chain
                        )
                    )
                )
            else:
                self.stdout.write(
                    self.style.WARNING(
                        "WARNING: Chain, pdb: {}, chain: {} doesn't exist".format(
163
164
                            pdb,
                            chain,
165
166
167
                        )
                    )
                )