gcp_publicload.py

#!/usr/bin/env python
"""
Uploads files to gcs buckets
Expectations:
1. Setup gcp credentials (via GOOGLE_APPLICATION_CREDENTIALS env var)
2. Setup SCRATCH_UPLOAD_BUCKET to the gcp bucket to place the files.

Examples:

# help
gcp_publicload -h

# upload gcp_publicload as <uuid>
gcp_publicload gcp_publicload

# upload gcp_publicload as <uuid>.py and set content type as plain text
gcp_publicload gcp_publicload -e py -ct text/plain

# upload gcp_publicload with the same name. Set the content type as plain text
gcp_publicload gcp_publicload -kn -ct text/plain

# (osx) upload a png image in clipboard (needs pasteboard package)
(osx) gcp_publicload
"""

import argparse
import os
import uuid
from pathlib import Path

import importlib

import platform
from google.cloud import storage

COMMON_CONTENT_TYPES = {
    'application/octet-stream',
    'application/pdf',
    'application/json',
    'application/xml',
    'application/zip',
    'audio/mpeg',
    'image/gif',
    'image/jpeg',
    'image/png',
    'image/tiff',
    'image/x-icon',
    'image/vnd.djvu',
    'image/svg+xml',
    'text/css',
    'text/csv',
    'text/html',
    'text/plain',
    'text/xml',
    'video/mpeg',
    'video/mp4',
    'video/x-flv',
    'video/webm'
}


def arg_parser():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('filename',
                        help='File to upload. '
                             'On osx, if this is not set, we read a png file from '
                             'clipboard and upload that', nargs='?')
    parser.add_argument('--keep_name', '-kn', action='store_true', default=False)
    parser.add_argument('--override', action='store_true', default=False)
    parser.add_argument(
        '--content_type', '-ct',
        help='Content type of the blob to upload.',
        type=str
    )
    parser.add_argument(
        '--extension', '-e',
        help=('Set this extension to the uploaded file. By default the extension will be '
              'preserved. To disable set this to _none_'),
        default='_preserve_',
        type=str
    )
    return parser


def fix_content_type(content_type):
    if not content_type:
        return None
    content_type = content_type.lower()
    if content_type not in COMMON_CONTENT_TYPES:
        maybe = input("Specified content type seems uncommon. Type yes to continue: ")
        if maybe.lower() != 'yes':
            print("Aborting")
            exit(1)
    return content_type


def fix_extension(filepath, extension):
    if extension == '_none_':
        return ''
    elif extension != '_preserve_':
        return '.' + extension

    name = os.path.basename(filepath)
    rsplit = name.rsplit('.', 1)
    if len(rsplit) == 1:  # The file name has no extension
        return ''
    else:
        return '.' + rsplit[-1]


def get_png_from_clipboard_osx():
    try:
        importlib.import_module('pasteboard')
    except Exception:
        print("filename not set. If you want to upload png from clipboard, install `pasteboard`")
        exit(1)
    import pasteboard
    pb = pasteboard.Pasteboard()
    png_data = pb.get_contents(type=pasteboard.PNG)
    if not png_data:
        arg_parser().error("filename not set and no png in clipboard")
    pngfile = f"/tmp/{uuid.uuid4()}.png"
    Path(pngfile).write_bytes(png_data)
    print("Uploading png file from clipboard")
    print(f"This is stored in {pngfile} temporarily")
    assert input("Type yes to continue. ").strip().lower() == 'yes'
    return pngfile


def upload_gcs(file_to_upload, content_type, extension, override):
    if not file_to_upload and platform.system() == 'Darwin':
        file_to_upload = get_png_from_clipboard_osx()
    content_type = fix_content_type(content_type)
    bucket = storage.Client().get_bucket(os.environ['SCRATCH_UPLOAD_BUCKET'])
    blobname = os.path.basename(file_to_upload)
    if not args.keep_name:
        extension = fix_extension(file_to_upload, extension)
        blobname = f'{uuid.uuid4()}{extension}'
    blob = bucket.blob(blobname)
    kwargs = {}
    if not override:
        # Ensure blob does not exist.
        kwargs = {'if_generation_match': 0}
    blob.metadata = {'orig': os.path.basename(file_to_upload)}
    blob.upload_from_filename(filename=file_to_upload, content_type=content_type, **kwargs)
    print(blob.public_url)
    print(f'gs://{blob.bucket.name}/{blob.name}')
    return blob


if __name__ == '__main__':
    args = arg_parser().parse_args()
    upload_gcs(args.filename, args.content_type, args.extension, args.override)