IniMotif/xmlparser.py at master · chengl7/IniMotif · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#!/usr/bin/env python3

import sys
import requests
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import fromstring, ElementTree
import re
import urllib.request
from itertools import islice

def getXML(search, dataType, **kwargs):
    # download an xml file for the specified search terms

    # if search query is entirely numeric, download by taxon
    if all(charac.isdigit() for charac in search):
        # ammend datatype for api
        if dataType == "run":
            dataType = "read_run"
        elif dataType == "study":
            dataType = "read_study"
        elif dataType == "experiment":
            dataType = "read_experiment"
        elif dataType == "sample":
            dataType = "read_sample"
        else:
            print ("Datatype is not recognised. Supported values are: run, study or experiment")
            exit()
        # build the url for the query and download the xml file
        build_url = {"accession": search,
                    "result": dataType
                    }
        response = requests.get("https://www.ebi.ac.uk/ena/browser/api/xml/links/taxon", params=build_url)
    # else use free text search
    else:
        # ammend datatype for api
        if dataType == "run":
            dataType = "sra-run"
        elif dataType == "study":
            dataType = "sra-study"
        elif dataType == "experiment":
            dataType = "sra-experiment"
        elif dataType == "sample":
            dataType = "sra-sample"
        else:
            print ("Datatype is not recognised. Supported values are: run, study or experiment")
            exit()
        # build the url for the query and download the xml file
        build_url = {"domain": dataType,
                    "query": search
                    }
        response = requests.get("https://www.ebi.ac.uk/ena/browser/api/xml/textsearch", params=build_url)

    # write to file
    with open('ena.xml', 'wb') as outfile:
        outfile.write(response.content)

def parseXMLgetFTP(xmlfile, dataType, numDown):
    # parse the xml file for http links which contain information on the fastq files
    # open the http links and write the result to file

    # create element tree object
    tree = ET.parse(xmlfile)
    # get root element
    root = tree.getroot()

    httplinks = []

    if numDown:
        # iterate xml file for http links, use number of downloads if applicable
        for item in islice(root.findall('.//XREF_LINK[DB="ENA-FASTQ-FILES"]'), 0, int(numDown)):
            httplinks.append(item.find('ID').text)
    else:
        for item in root.findall('.//XREF_LINK[DB="ENA-FASTQ-FILES"]'):
            httplinks.append(item.find('ID').text)

    # fetch http data and write to file
    with open('fastq.txt', 'wb') as outfile:
        for url in httplinks:
            response = requests.get(url)
            outfile.write(response.content)

def parseFTPgetFASTQ(ftpinfo):
    # parse the txt file with the fastq info for the ftp links and download

    # use regex to compile ftp links
    regexFTP = re.compile("ftp.")
    # use regex to compile filesizes
    regexSize = re.compile(r"\d*;\d*|\d")

    # collate filesizes, filenames, ftp links, and sequencing type
    fileSize = []
    filename = []
    ftplink = []
    seqType = []

    with open(ftpinfo, 'r') as infile:
        for line in infile:
            # collate all the filesizes
            try:
                linesplit = line.split()[3]
            except IndexError:
                linesplit = "null"
            if regexSize.match(linesplit):
            # check for paired fastq files
                for elem in linesplit.split(";", 2):
                    fileSize.append(elem)
            # collate filenames and ftplinks
            try:
                linesplit = line.split()[1]
            except IndexError:
                linesplit = "null"
            if regexFTP.match(linesplit):
            # check for paired fastq files
                if len(linesplit.split(";", 2)) >= 2:
                    seqType.append("PAIRED")
                else:
                    seqType.append("SINGLE")
                for elem in linesplit.split(";", 2):
                    filename.append(elem[elem.rfind("/")+1:])
                    ftplink.append("ftp://" + elem)

    # sum total filesizes and print to terminal
    add = [int(x) for x in fileSize]
    tot = sum(add)/10**9
    print("You are about to download " + str(round(tot, 2)) +  " GB of files")
    sys.stdout.flush()

    # fetch fastqs
    for link, name in zip(ftplink, filename):
        urllib.request.urlretrieve(link, name)

    return seqType