Music-Downloader/helper/sites/request.py

import requests
import urllib.parse

from helper import console


class Request:
    def __init__(self, site):
        self.site = site
        self.headers = {}
        self.cookie = {}

    def add_header(self, key, value):
        self.headers[key] = value
        return self

    def add_cookie(self, key, value):
        self.cookie[key] = value
        return self

    def search(self, keywords):
        if self.site.query.encoding == '+':
            query = self.site.query.format_query(urllib.parse.quote(keywords.replace(' ', '+')))
        else:
            query = self.site.query.format_query(urllib.parse.quote(keywords))

        console.output('Request method: {0}'.format(self.site.query.method), level=console.DBG_INFO)

        if self.site.query.method == 'GET':
            request = requests.get(
                '{0}?{1}'.format(self.site.query.url_path, query),
                headers=self.headers,
                cookies=self.cookie,
                allow_redirects=True)
        elif self.site.query.method == 'POST':
            request = requests.post(
                self.site.query.url_path,
                data=query,
                headers=self.headers,
                cookies=self.cookie,
                allow_redirects=True)
        else:  # FLAT
            request = requests.get(
                '{0}/{1}'.format(self.site.query.url_path, query),
                headers=self.headers,
                cookies=self.cookie,
                allow_redirects=True)

        console.output('Request response: {0} {1}'.format(request.status_code, request.reason), level=console.DBG_INFO)

        return self.site.structure.parse(request.text)

    def parse_url(self, url):
        o = urllib.parse.urlparse(url)

        if o.query != '':
            o.query = urllib.parse.parse_qs(o.query)

        return o


class RawRequest:
    @staticmethod
    def get(link, parameters=None, headers=None, cookies=None, stream=False):
        request = requests.get(
            link,
            data=parameters,
            headers=headers,
            cookies=cookies,
            allow_redirects=True,
            stream=stream)
        if request.status_code == 200:
            return request

    @staticmethod
    def post(link, parameters=None, headers=None, cookies=None, stream=False):
        request = requests.post(
            link,
            data=parameters,
            headers=headers,
            cookies=cookies,
            allow_redirects=True,
            stream=stream)
        if request.status_code == 200:
            return request