<commit_before>"""Download and install structured genome data and aligner index files.

Downloads prepared FASTA, indexes for aligners like BWA, Bowtie and novoalign
and other genome data in automated pipelines. Specify the genomes and aligners
to use in an input biodata.yaml configuration file.

The main targets are fabric functions:

  - install_data -- Install biological data from scratch, including indexing genomes.
  - install_data_s3 -- Install biological data, downloading pre-computed indexes from S3.
  - upload_s3 -- Upload created indexes to biodata S3 bucket.

"""
import os
import operator
import socket
import subprocess
from contextlib import contextmanager

from fabric.api import *
from fabric.contrib.files import *
from fabric.context_managers import path
try:
    import yaml
except ImportError:
    yaml = None
try:
    import boto
except ImportError:
    boto = None

from cloudbio.biodata import galaxy
from cloudbio.biodata.dbsnp import download_dbsnp
from cloudbio.biodata.rnaseq import download_transcripts
from cloudbio.custom import shared

# -- Configuration for genomes to download and prepare

class _DownloadHelper:
    def __init__(self):
        self.config = {}

    def ucsc_name(self):
        return None

    def _exists(self, fname, seq_dir):
        """Check if a file exists in either download or final destination.
        """
        return env.safe_exists(fname) or env.safe_exists(os.path.join(seq_dir, fname))

class UCSCGenome(_DownloadHelper):
    def __init__(self, genome_name, dl_name=None):
        _DownloadHelper.__init__(self)
        self.data_source = "UCSC"
        self._name = genome_name
        self.dl_name = dl_name if dl_name is not None else genome_name
        self._url = "ftp://hgdownload.cse.ucsc.edu/goldenPath/%s/bigZips" % \
                genome_name

    def ucsc_name(self):
        return self._name

    def _karyotype_sort(self, xs):
        """Sort reads in karyotypic order to work with GATK's defaults.
        """
        def karyotype_keyfn(x):
            base = os.path.splitext(os.path.basename(x))[0]
            if base.startswith("chr"):
                base = base[3:]
            parts = base.split("_")
            try:
                parts[0] =  int(parts[0])
            except ValueError:
                pass
            # unplaced at the very end
            if parts[0] == "Un":
                parts.insert(0, "z")
            # mitochondrial special case -- after X/Y
            elif parts[0] in ["M", "MT"]:
                parts.insert(0, "x")
            # sort random and extra chromosomes after M
            elif len(parts) > 1:
                parts.insert(0, "y")
            return parts
        return sorted(xs, key=karyotype_keyfn)

    def download(self, seq_dir):
        zipped_file = None
        genome_file = "%s.fa" % self._name
        if not self._exists(genome_file, seq_dir):
            zipped_file = self._download_zip(seq_dir)
            if zipped_file.endswith(".tar.gz"):
                env.safe_run("tar -xzpf %s" % zipped_file)
            elif zipped_file.endswith(".zip"):
                env.safe_run("unzip %s" % zipped_file)
            elif zipped_file.endswith(".gz"):
                env.safe_run("gunzip -c %s > out.fa" % zipped_file)
            else:
                raise ValueError("Do not know how to handle: %s" % zipped_file)
            tmp_file = genome_file.replace(".fa", ".txt")
            result = env.safe_run_output("find . -name '*.fa'")
            result = self._karyotype_sort([x.strip() for x in result.split("\n")])
            env.safe_run("cat %s > %s" % (" ".join(result), tmp_file))
            env.safe_run("rm -f *.fa")
            env.safe_run("mv %s %s" % (tmp_file, genome_file))
        return genome_file, [zipped_file]

    def _download_zip(self, seq_dir):<commit_msg>Fix for karyotype sorting multifasta genomes.<commit_after>"""Download and install structured genome data and aligner index files.

Downloads prepared FASTA, indexes for aligners like BWA, Bowtie and novoalign
and other genome data in automated pipelines. Specify the genomes and aligners
to use in an input biodata.yaml configuration file.

The main targets are fabric functions:

  - install_data -- Install biological data from scratch, including indexing genomes.
  - install_data_s3 -- Install biological data, downloading pre-computed indexes from S3.
  - upload_s3 -- Upload created indexes to biodata S3 bucket.

"""
import os
import operator
import socket
import subprocess
from contextlib import contextmanager

from fabric.api import *
from fabric.contrib.files import *
from fabric.context_managers import path
try:
    import yaml
except ImportError:
    yaml = None
try:
    import boto
except ImportError:
    boto = None

from cloudbio.biodata import galaxy
from cloudbio.biodata.dbsnp import download_dbsnp
from cloudbio.biodata.rnaseq import download_transcripts
from cloudbio.custom import shared

# -- Configuration for genomes to download and prepare

class _DownloadHelper:
    def __init__(self):
        self.config = {}

    def ucsc_name(self):
        return None

    def _exists(self, fname, seq_dir):
        """Check if a file exists in either download or final destination.
        """
        return env.safe_exists(fname) or env.safe_exists(os.path.join(seq_dir, fname))

class UCSCGenome(_DownloadHelper):
    def __init__(self, genome_name, dl_name=None):
        _DownloadHelper.__init__(self)
        self.data_source = "UCSC"
        self._name = genome_name
        self.dl_name = dl_name if dl_name is not None else genome_name
        self._url = "ftp://hgdownload.cse.ucsc.edu/goldenPath/%s/bigZips" % \
                genome_name

    def ucsc_name(self):
        return self._name

    def _karyotype_sort(self, xs):
        """Sort reads in karyotypic order to work with GATK's defaults.
        """
        def karyotype_keyfn(x):
            base = os.path.splitext(os.path.basename(x))[0]
            if base.startswith("chr"):
                base = base[3:]
            parts = base.split("_")
            try:
                parts[0] =  int(parts[0])
            except ValueError:
                pass
            # unplaced at the very end
            if parts[0] == "Un":
                parts.insert(0, "z")
            # mitochondrial special case -- after X/Y
            elif parts[0] in ["M", "MT"]:
                parts.insert(0, "x")
            # sort random and extra chromosomes after M
            elif len(parts) > 1:
                parts.insert(0, "y")
            return parts
        return sorted(xs, key=karyotype_keyfn)

    def _split_multifasta(self, fasta_file):
        chrom = ""
        file_handle = None
        file_names = []
        out_dir = os.path.dirname(fasta_file)
        with open(fasta_file) as in_handle:
            for line in in_handle:
                if line.startswith(">"):
                    chrom = line.split(">")[1].strip()
                    file_handle.close() if file_handle else None
                    file_names.append(chrom + ".fa")
                    file_handle = open(os.path.join(out_dir, chrom + ".fa"), "w")
                    file_handle.write(line)
                else:
                    file_handle.write(line)
        file_handle.close()
        return file_names


    def download(self, seq_dir):
        zipped_file = None
        genome_file = "%s.fa" % self._name
        if not self._exists(genome_file, seq_dir):
            zipped_file = self._download_zip(seq_dir)
            if zipped_file.endswith(".tar.gz"):
                env.safe_run("tar -xzpf %s" % zipped_file)
            elif zipped_file.endswith(".zip"):
                env.safe_run("unzip %s" % zipped_file)
            elif zipped_file.endswith(".gz"):
                env.safe_run("gunzip -c %s > out.fa" % zipped_file)
            else:
                raise ValueError("Do not know how to handle: %s" % zipped_file)
            tmp_file = genome_file.replace(".fa", ".txt")
            result = env.safe_run_output("find `pwd` -name '*.fa'")
            result = [x.strip() for x in result.split("\n")]
            if len(result) == 1:
                result = self._split_multifasta(result[0])
            result = self._karyotype_sort(result)
            env.safe_run("cat %s > %s" % (" ".join(result), tmp_file))
            env.safe_run("rm -f *.fa")
            env.safe_run("mv %s %s" % (tmp_file, genome_file))
        return genome_file, [zipped_file]

    def _download_zip(self, seq_dir):
{'fraction_non_alphanumeric': 0.0997197571228398, 'fraction_numerical': 0.003619803829985988, 'mean_word_length': 2.2369614512471654, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6930270', 'n_tokens_mistral': 2641, 'n_tokens_neox': 2501, 'n_words': 711}