PrivacyScanData/S3/NewFind/s3transfer-develop/scripts/performance/download-file

#!/usr/bin/env python
"""
Downloads a file using s3transfer. You can also chose how type of file that
is downloaded (i.e. filename, seekable, nonseekable).

Usage
=====

NOTE: Make sure you run ``pip install -r requirements-dev.txt`` before running.

To download a file::

    ./download-file --file-name myfilename --file-type filename \\
        --s3-bucket mybucket --s3-key mykey

"""
import argparse

from botocore.session import get_session

from s3transfer.manager import TransferManager


class NonSeekableWriter:
    """A wrapper to hide the ability to seek for a fileobj"""

    def __init__(self, fileobj):
        self._fileobj = fileobj

    def write(self, b):
        return self._fileobj.write(b)


class Downloader:
    def download(self, args):
        session = get_session()
        client = session.create_client('s3')
        file_type = args.file_type
        if args.debug:
            session.set_debug_logger('')
        with TransferManager(client) as manager:
            getattr(self, 'download_' + file_type)(
                manager, args.file_name, args.s3_bucket, args.s3_key
            )

    def download_filename(self, manager, filename, bucket, s3_key):
        manager.download(bucket, s3_key, filename)

    def download_seekable(self, manager, filename, bucket, s3_key):
        with open(filename, 'wb') as f:
            future = manager.download(bucket, s3_key, f)
            future.result()

    def download_nonseekable(self, manager, filename, bucket, s3_key):
        with open(filename, 'wb') as f:
            future = manager.download(bucket, s3_key, NonSeekableWriter(f))
            future.result()


def main():
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument('--file-name', required=True, help='The name of file')
    parser.add_argument(
        '--file-type',
        choices=['filename', 'seekable', 'nonseekable'],
        required=True,
        help='The way to represent the file when downloading',
    )
    parser.add_argument(
        '--s3-bucket',
        required=True,
        help='The S3 bucket to download the file to',
    )
    parser.add_argument(
        '--s3-key', required=True, help='The key to download to'
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        help='Whether to turn debugging on. This will get printed to stderr',
    )
    args = parser.parse_args()
    Downloader().download(args)


if __name__ == '__main__':
    main()