Python Google Drive Downloader

gdrive_downloader.py

__all__ = ['gdrive_downloader']

import requests
import re
import os
import warnings


_codefinder = re.compile(r'confirm=(.+)&id=')
_namefinder = re.compile(r'filename="(.+)";')
    
def gdrive_downloader(id, outpath=None, overwrite=False, chunksize=8192, verbose=0):
    """
    Arguments
    ---------
    id : str
        Unique google id
    outpath : str or None
        if None, the file's content disposition will be used.
    overwrite : bool
        If True and the file exists, it will be overwritten.
    chunksize : int
        Size of chunks to write from the incoming stream to file at a time.
    Returns
    -------
    outpath : str
    """
    # used like this
    # download_google <id> <name of item.extension>
    session = requests.Session()
    url1 = f"https://drive.google.com/uc?export=download&id={id}"
    if verbose > 0:
        print(url1, flush=True)
    
    ir = session.get(url1)
    if verbose > 0:
        print(ir.text, flush=True)

    results = _codefinder.findall(ir.text)
    if len(results) == 0:
        print(ir.text)
        raise ValueError('Unable to get confirmation code to download; likely too many downloads.')

    url2 = f"https://drive.google.com/uc?export=download&confirm={results[0]}&id={id}"
    with session.get(url2, stream=True) as fr:
        if verbose == 1:
            print(fr.headers['content-disposition'], flush=True)
        elif verbose > 1:
            print(fr.headers, flush=True)

        results = _namefinder.findall(fr.headers.get('content-disposition', ''))
        if outpath is None:
            outpath = results[0]
        else:
            raise KeyError(f'{results} and outpath is None')

        fr.raise_for_status()
        if os.path.exists(outpath) and not overwrite:
            warnings.warn(f'{outpath} exists; set overwrite=True to continue')
        else:
            if verbose > 0:
                print('Writing')
            with open(outpath, 'wb') as f:
                for chunk in fr.iter_content(chunk_size=chunksize): 
                    f.write(chunk)

    return outpath


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    aa = parser.add_argument
    aa('--verbose', action='count', default=0, help='Add -v to increase verbosity')
    aa('--chunksize', type=int, default=8192, help='Overwrite if outpath exists')
    aa('--overwrite', type=bool, default=False, action='store_true', help='Overwrite if outpath exists')
    aa('--outpath', type=str, default=None, help='Path to save; defaults to filename')
    aa('id', type=str)
    
    args = parser.parse_args()
    gdrive_downloader(**vars(args))