utils.py

"""
This module contains a set of functions to be used by others.
Some of them are ripped from https://github.com/coursera-dl/
"""

import time
import math
import sys
import os
import re
import errno


class DownloadProgress(object):
    """
    Report download progress.
    Ripped from https://github.com/coursera-dl/
    """

    def __init__(self, start, total):
        if total in [0, '0', None]:
            self._total = None
        else:
            self._total = int(total)

        self._current = int(start)
        self._start = int(start)
        self._time_start = 0
        self._time_now = 0

        self._finished = False

    def start(self):
        # To set _time_now & _time_start
        self._time_now = time.time()
        self._time_start = self._time_now

    def stop(self):
        # To set _time_now & _finished
        self._time_now = time.time()
        self._finished = True
        # If total is not set, set _curent as _total
        if self._total is None:
            self._total = self._current
        self.report_progress()
        # If not curent==total
        if self._total != self._current:
            raise Exception('Error: Stopped abnormally.')

    def read(self, bytes):
        # Update time
        self._time_now = time.time()
        # Update current
        self._current += bytes

        self.report_progress()

    def calc_percent(self):
        if self._total is None:
            return '--%'
        percentage = int(float(self._current) / float(self._total) * 100.0)
        done = int(percentage / 2)
        return '[{0: <50}] {1}%'.format(done * '>', percentage)

    def calc_speed(self):
        dif = self._time_now - self._time_start
        if self._current == 0 or float(dif) < 0.001:  # One millisecond
            return '---b/s    '
        return '{0}/s    '.format(format_bytes(float(self._current - self._start) / dif))

    def report_progress(self):
        """
        Report download progress.
        Ripped from https://github.com/coursera-dl/
        """
        percent = self.calc_percent()
        total = format_bytes(self._total)

        speed = self.calc_speed()
        total_speed_report = '{0} at {1}'.format(total, speed)

        report = '\r{0: <56} {1: >20}'.format(percent, total_speed_report)

        if self._finished:
            print(report)
        else:
            print(report + "\r"),

        sys.stdout.flush()


def format_bytes(bytes):
    """
    Get human readable version of given bytes.
    Ripped from https://github.com/rg3/youtube-dl
    """
    if bytes is None:
        return 'N/A'
    if type(bytes) is str:
        bytes = float(bytes)
    if bytes == 0.0:
        exponent = 0
    else:
        exponent = int(math.log(bytes, 1024.0))
    suffix = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'][exponent]
    converted = float(bytes) / float(1024 ** exponent)
    return '{0:.2f}{1}'.format(converted, suffix)


def mkdir_p(path, mode=0o777):
    """
    Create subdirectory hierarchy given in the paths argument.
    Ripped from https://github.com/coursera-dl/
    """

    try:
        os.makedirs(path, mode)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise


def download_file(session, url, filename):
    attempts_count = 0
    error_msg = ''

    while attempts_count < 2:

        r = session.get(url, stream=True)
        # if not succeed
        if r.status_code is not 200:
            if r.reason:
                error_msg = r.reason + ' ' + str(r.status_code)
            else:
                error_msg = 'HTTP Error ' + str(r.status_code)

            if attempts_count + 1 < 2:
                wait_interval = 2 ** (attempts_count + 1)
                msg = 'Error downloading, will retry in {0} seconds ...'
                print(msg.format(wait_interval))
                time.sleep(wait_interval)
                attempts_count += 1
                continue
            else:
                raise Exception('Connection Error: %s' % error_msg)

        content_length = r.headers.get('content-length')
        # Create DownloadProcess here.       args: start total
        progress = DownloadProgress(0, content_length)
        chunk_sz = 1048576

        with open(filename, 'wb') as f:
            progress.start()
            while True:
                data = r.raw.read(chunk_sz, decode_content=True)
                if not data:
                    progress.stop()
                    break
                progress.read(len(data))
                f.write(data)
        r.close()
        break


def resume_download_file(session, url, filename, overwrite=False):
    if os.path.exists(filename) and not overwrite:
        resume_len = os.path.getsize(filename)
        file_mode = 'ab'
    else:
        resume_len = 0
        file_mode = 'wb'

    attempts_count = 0
    error_msg = ''

    while attempts_count < 2:

        session.headers['Range'] = 'bytes=0-'
        r = session.get(url, stream=True)

        if r.status_code != 200 and r.status_code != 206:
            if r.reason:
                error_msg = r.reason + ' ' + str(r.status_code)
            else:
                error_msg = 'HTTP Error ' + str(r.status_code)

            if attempts_count + 1 < 2:
                wait_interval = 2 ** (attempts_count + 1)
                msg = 'Error downloading, will retry in {0} seconds ...'
                print(msg.format(wait_interval))
                time.sleep(wait_interval)
                attempts_count += 1
                continue
            else:
                raise Exception('Connection Error: %s' % error_msg)

        total_length = r.headers.get('content-length')

        if resume_len != 0:
            if total_length is None or resume_len == int(total_length):
                print('Already downloaded.')
                break

        session.headers['Range'] = 'bytes=%d-' % (resume_len)
        r = session.get(url, stream=True)

        progress = DownloadProgress(resume_len, total_length)
        chunk_sz = 1048576

        with open(filename, file_mode) as f:
            progress.start()
            while True:
                data = r.raw.read(chunk_sz, decode_content=True)
                if not data:
                    progress.stop()
                    break
                progress.read(len(data))
                f.write(data)
        r.close()
        break


def clean_filename(s):
    """
    Amend  string to legal filename
    """
    s = s.strip()\
        .replace(':', '_') \
        .replace('/', '_') \
        .replace('\x00', '_') \
        .replace('\xa0', '') \
        .replace(' ', '_')
    s = re.sub('[\n\\\*><\?\"\|\t\s]', '', s)
    return s