exportsvg.py

import math
import os
import re
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
from progress.bar import IncrementalBar
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from .Message import Message


class DownloadException(Exception):
    pass

session = None
COLLECTIONS_API_URL = 'https://api.svgrepo.com/collections/'
COLLECTION_API_URL = 'https://api.svgrepo.com/collection/'
COLLECTIONS_PAGE_SIZE = 100
COLLECTION_ITEMS_PAGE_SIZE = 50
DOWNLOAD_WORKERS = int(os.environ.get('SVGREPODL_WORKERS', '8'))


def _new_session():
    retry = Retry(total=2, backoff_factor=1)
    current_session = requests.Session()
    current_session.mount('https://', HTTPAdapter(max_retries=retry))
    return current_session


def _fetch_json(url, *, params=None):
    response = session.get(url, params=params, timeout=30)
    if response.status_code != 200:
        raise DownloadException()
    return response.json()


def list_collections(category='all'):
    del category

    global session
    session = _new_session()

    start = 0
    total_pages = None
    while True:
        payload = _fetch_json(
            COLLECTIONS_API_URL,
            params={'limit': COLLECTIONS_PAGE_SIZE, 'start': start},
        )
        collections = payload.get('collections', [])
        if not collections:
            break

        if total_pages is None:
            count = int(payload['count'])
            limit = int(payload['limit'])
            total_pages = math.ceil(count / limit)

        limit = int(payload['limit'])
        current_page = (start // limit) + 1
        print(f'page {current_page}/{total_pages}', file=sys.stderr)
        for item in collections:
            print(f'https://www.svgrepo.com/collection/{item["slug"]}/')

        start += len(collections)


def _download_item(link, path):
    aid = os.path.basename(os.path.dirname(link))
    dest = os.path.join(path, aid + '-' + os.path.basename(link))
    if os.path.exists(dest):
        return 'skipped'

    response = session.get(link, timeout=30)
    if response.headers.get('content-type') != 'image/svg+xml':
        print("err", link, file=sys.stderr)
        return 'error'

    with open(dest, 'wb') as handle:
        handle.write(response.content)
    return 'downloaded'


def download_items(all_links, path, bar):
    with ThreadPoolExecutor(max_workers=DOWNLOAD_WORKERS) as executor:
        futures = [executor.submit(_download_item, link, path) for link in all_links]
        for future in as_completed(futures):
            future.result()
            bar.next()


def downloader(url, path, only_list=False, collection=''):
    """
    Download a collection (or a search)

    Arguments:
    url {[string]} -- URL of SVGREPO Collection
    """
    is_search = '/vectors/' in url
    if is_search:
        raise DownloadException()

    global session
    session = _new_session()

    os.makedirs(path, exist_ok=True)

    collection_slug = collection or re.match(r'.*/collection/([^/?#&]+)', url).group(1)
    start = 0
    total_pages = None
    while True:
        payload = _fetch_json(
            COLLECTION_API_URL,
            params={
                'term': collection_slug,
                'limit': COLLECTION_ITEMS_PAGE_SIZE,
                'start': start,
            },
        )
        icons = payload.get('icons', [])
        if not icons:
            if start == 0:
                raise DownloadException()
            break

        if total_pages is None:
            total_pages = math.ceil(int(payload['count']) / int(payload['limit']))

        all_links = [
            f'https://www.svgrepo.com/show/{icon["id"]}/{icon["slug"]}.svg'
            for icon in icons
        ]

        if only_list:
            print("\n".join([collection_slug + "\t" + e for e in all_links]))
        else:
            limit = int(payload['limit'])
            current_page = (start // limit) + 1
            bar = IncrementalBar('📥 Icons URLs page %d/%d' % (current_page, total_pages), max=len(all_links))
            download_items(all_links, path, bar)
            bar.finish()

        start += len(icons)
    if not only_list:
        Message.success('🎉 Finished')