update_mrc_files.py

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
from celery import states
import os
import sys
import glob


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


def isolevel_float(arg):
    """ Type function for argparse - a float within some predefined bounds """
    try:
        f = float(arg)
    except ValueError:
        raise CommandError("Must be a floating point number")
    if f < 0 or f > 1.0:
        raise CommandError("Argument must be < 1.0 and > 0")
    return f


class Command(AppCommand):
    """
    Command to update mrc files
    """

    help = "liste all mrc files in a directory and update files in the database"

    def __init__(
        self,
        stdout=None,
        stderr=None,
        no_color=False,
        force_color=False,
        task=None,
    ):
        super(Command, self).__init__(
            stdout=stdout, stderr=stderr, no_color=no_color, force_color=force_color
        )
        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)

    def add_arguments(self, parser):
        parser.add_argument(
            "-mp",
            "--mrcpath",
            type=dir_path,
            help="root directory with all mrc files",
        )
        parser.add_argument(
            "-pa",
            "--pattern",
            dest="pattern",
            help="pattern for select files",
            default="*.mrc",
        )
        parser.add_argument(
            "-t",
            "--type",
            dest="type",
            choices=["druggability", "interactibility"],
            default="druggability",
            help="files type, druggability or interactibility",
        )
        parser.add_argument(
            "-l",
            "--label",
            dest="label",
            help="label for files, used as label in the interface",
        )
        parser.add_argument(
            "-di",
            "--default_isolevel",
            dest="default_isolevel",
            help="isolevel default value",
            default=0.7,
            type=isolevel_float,
        )

    def handle(self, *args, **options):
        """
        Perform the command's actions
        """
        typefile = options["type"]
        pattern = options["pattern"]
        mrcpath = options["mrcpath"]
        label = options.get("label")
        if typefile == "interactibility" and not label:
            raise CommandError("Error: Interactibility need a label")
        if typefile == "druggability":
            storage_path = Chain.mrc_file.field.storage.path("content/mrc_files")
        else:
            storage_path = InteractFile.interact_file.field.storage.path(
                "content/interact_files"
            )
        if not os.path.exists(storage_path):
            os.makedirs(storage_path)
        list_files = glob.glob(os.path.join(mrcpath, pattern))
        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:
                chain = splitname[-1].split("_")[0]
                pdb = splitname[0]
            elif len(splitname) == 3:
                chain = splitname[1].split("_")[0]
                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))
                        dbchains[0].default_isolevel = options["default_isolevel"]
                        dbchains[0].save()
                else:
                    interactfile, created = InteractFile.objects.get_or_create(
                        chain=dbchains[0],
                        label=label,
                    )
                    with open(data, "rb") as f:
                        interactfile.interact_file.delete()
                        interactfile.interact_file.save(os.path.basename(data), File(f))
                    interactfile.default_isolevel = options["default_isolevel"]
                    interactfile.save()
                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(
                            pdb,
                            chain,
                        )
                    )
                )