# Path: ob_pipelines/s3.py
# def s3args(rm_local_outpath=False):
#     """Sync S3 path arguments with behind-the-scenes S3 transfers
# 
#     When decorating a function, s3args downloads all arguments that 
#     look like S3 paths to temporary files and swaps the local temp
#     filepath as the new argument. If the S3 path does not exist, it 
#     is assumed to be an output, and s3args uploads the tempfile 
#     back to S3 after the command is complete.
# 
#     This works great with Luigi, which checks for existence of inputs
#     and non-existence of outputs before running a Task.
# 
#     Keyword args are passed directly without syncing, for now.
#     """
#     def s3args_decorator(f):
#         @wraps(f)
#         def local_fn(*args, **kwargs):
# 
#             # Swap the S3 path arguments for local temporary files/folders
#             local_args, s3_downloads, s3_uploads = swap_args(args, rm_local_outpath=rm_local_outpath)
# 
#             try:
#                 # Download inputs
#                 logger.info('syncing from S3')
#                 for s3_path, local_path in s3_downloads.items():
#                     download_file_or_folder(s3_path, local_path)
# 
#                 # Run command and save output
#                 out = f(*local_args, **kwargs)
# 
#                 # Upload outputs
#                 logger.info('uploading to S3')
#                 for s3_path, local_path in s3_uploads.items():
#                     upload_file_or_folder(s3_path, local_path)
#             finally:
#                 # Remove local files
#                 local_paths = list(s3_downloads.values()) + list(s3_uploads.values())
#                 for local_path in local_paths:
#                     remove_file_or_folder(local_path)
# 
#             return out
# 
#         return local_fn
#     return s3args_decorator
# 
# def key_exists(bucket, key):
#     """Check for existence of S3 key"""
#     try:
#         if key.endswith('/'):
#             return 'Contents' in s3.list_objects(Bucket=bucket, Prefix=key)
#         else:
#             s3_resource.Object(bucket, key).load()
#             return True
#     except botocore.exceptions.ClientError as e:
#         return False
# 
# def path_to_bucket_and_key(path):
#     (scheme, netloc, path, params, query, fragment) = urlparse(path)
#     path_without_initial_slash = path[1:]
#     return netloc, path_without_initial_slash
# 
# def create_tmp_from_key(key): 
#     """Create local temp file or folder depending on key path"""
#     if key.endswith('/'):
#         local_tmp = mkdtemp(
#             prefix=op.basename(key.rstrip('/')) + '_', 
#             dir=SCRATCH_DIR)
#         local_tmp = local_tmp if local_tmp.endswith('/') else local_tmp + '/'
#     else:
#         fname = op.basename(key)
#         base, ext = op.splitext(fname)
#         _, local_tmp = mkstemp(prefix=base + '_',
#                                suffix=ext,
#                                dir=SCRATCH_DIR)
#     return local_tmp
# 
# def upload_folder(folder, bucket, prefix):
#     # Upload directory
#     for fname in os.listdir(folder):
#         fpath = op.join(folder, fname)
#         key = prefix + fname
#         s3.upload_file(fpath, bucket, key,
#             ExtraArgs={'ServerSideEncryption': 'AES256'})
# 
# def download_folder(bucket, prefix, folder):
#     response = s3.list_objects(
#         Bucket=bucket, 
#         Prefix=prefix,
#         Delimiter='/'
#     )
#     if 'Contents' not in response and response['HTTPStatusCode'] == 200:
#         raise botocore.exceptions.ClientError
#     for key_dict in response['Contents']:
#         key = key_dict['Key']
#         fpath = op.join(folder, op.basename(key))
#         s3.download_file(bucket, key, fpath)
# 
# def swap_args(args, rm_local_outpath=False):
#     """Swap S3 paths in arguments with local paths
#     
#     If the S3 path exists, it's an input, download first and swap the arg
#     with a temporary filepath. Otherwise, it's an output, save for upload 
#     after the command.
# 
#     Returns: 
#         tuple of (local_args, s3_downloads, s3_uploads)
# 
#         where new_args contains the new argument list with local paths, 
#         and s3_outputs is a dict mapping local filepaths to s3 paths to 
#         transfer after execution.
#     """
#     s3_uploads = {}
#     s3_downloads = {}
#     local_args = []
#     for arg in args:
#         if not arg.startswith('s3://'):
#             local_args.append(arg)
#             continue
#         
#         src_bucket, src_key = path_to_bucket_and_key(arg)
#         local_tmp = create_tmp_from_key(src_key)
#         
#         # If key exists, add path to downloads, otherwise add path to 
#         # uploads and remove the file or folder
#         if key_exists(src_bucket, src_key):
#             s3_downloads[arg] = local_tmp
#         else:
#             s3_uploads[arg] = local_tmp
#             if rm_local_outpath:
#                 remove_file_or_folder(local_tmp)
# 
#         local_args.append(local_tmp)
# 
#     return local_args, s3_downloads, s3_uploads


import os
import os.path as op
import shutil
import boto3
import pytest
from subprocess import check_output
from tempfile import mkstemp, mkdtemp
from ob_pipelines.s3 import (
    s3args, key_exists, path_to_bucket_and_key, create_tmp_from_key,
    upload_folder, download_folder, swap_args
)


TEST_BUCKET = os.environ.get('TEST_BUCKET')
OUT_KEY = 'test_out'
IN_KEY = 'test_in'
IN_PREFIX = 'test_prefix/'
OUT_PREFIX = 'test_out_prefix/'
s3 = boto3.client('s3')


@pytest.fixture
def s3_outpath():
    s3.delete_object(Bucket=TEST_BUCKET, Key=OUT_KEY)
    return 's3://{}/{}'.format(TEST_BUCKET, OUT_KEY)



# Path: ob_pipelines/s3.py
# def s3args(rm_local_outpath=False):
#     """Sync S3 path arguments with behind-the-scenes S3 transfers
# 
#     When decorating a function, s3args downloads all arguments that 
#     look like S3 paths to temporary files and swaps the local temp
#     filepath as the new argument. If the S3 path does not exist, it 
#     is assumed to be an output, and s3args uploads the tempfile 
#     back to S3 after the command is complete.
# 
#     This works great with Luigi, which checks for existence of inputs
#     and non-existence of outputs before running a Task.
# 
#     Keyword args are passed directly without syncing, for now.
#     """
#     def s3args_decorator(f):
#         @wraps(f)
#         def local_fn(*args, **kwargs):
# 
#             # Swap the S3 path arguments for local temporary files/folders
#             local_args, s3_downloads, s3_uploads = swap_args(args, rm_local_outpath=rm_local_outpath)
# 
#             try:
#                 # Download inputs
#                 logger.info('syncing from S3')
#                 for s3_path, local_path in s3_downloads.items():
#                     download_file_or_folder(s3_path, local_path)
# 
#                 # Run command and save output
#                 out = f(*local_args, **kwargs)
# 
#                 # Upload outputs
#                 logger.info('uploading to S3')
#                 for s3_path, local_path in s3_uploads.items():
#                     upload_file_or_folder(s3_path, local_path)
#             finally:
#                 # Remove local files
#                 local_paths = list(s3_downloads.values()) + list(s3_uploads.values())
#                 for local_path in local_paths:
#                     remove_file_or_folder(local_path)
# 
#             return out
# 
#         return local_fn
#     return s3args_decorator
# 
# def key_exists(bucket, key):
#     """Check for existence of S3 key"""
#     try:
#         if key.endswith('/'):
#             return 'Contents' in s3.list_objects(Bucket=bucket, Prefix=key)
#         else:
#             s3_resource.Object(bucket, key).load()
#             return True
#     except botocore.exceptions.ClientError as e:
#         return False
# 
# def path_to_bucket_and_key(path):
#     (scheme, netloc, path, params, query, fragment) = urlparse(path)
#     path_without_initial_slash = path[1:]
#     return netloc, path_without_initial_slash
# 
# def create_tmp_from_key(key): 
#     """Create local temp file or folder depending on key path"""
#     if key.endswith('/'):
#         local_tmp = mkdtemp(
#             prefix=op.basename(key.rstrip('/')) + '_', 
#             dir=SCRATCH_DIR)
#         local_tmp = local_tmp if local_tmp.endswith('/') else local_tmp + '/'
#     else:
#         fname = op.basename(key)
#         base, ext = op.splitext(fname)
#         _, local_tmp = mkstemp(prefix=base + '_',
#                                suffix=ext,
#                                dir=SCRATCH_DIR)
#     return local_tmp
# 
# def upload_folder(folder, bucket, prefix):
#     # Upload directory
#     for fname in os.listdir(folder):
#         fpath = op.join(folder, fname)
#         key = prefix + fname
#         s3.upload_file(fpath, bucket, key,
#             ExtraArgs={'ServerSideEncryption': 'AES256'})
# 
# def download_folder(bucket, prefix, folder):
#     response = s3.list_objects(
#         Bucket=bucket, 
#         Prefix=prefix,
#         Delimiter='/'
#     )
#     if 'Contents' not in response and response['HTTPStatusCode'] == 200:
#         raise botocore.exceptions.ClientError
#     for key_dict in response['Contents']:
#         key = key_dict['Key']
#         fpath = op.join(folder, op.basename(key))
#         s3.download_file(bucket, key, fpath)
# 
# def swap_args(args, rm_local_outpath=False):
#     """Swap S3 paths in arguments with local paths
#     
#     If the S3 path exists, it's an input, download first and swap the arg
#     with a temporary filepath. Otherwise, it's an output, save for upload 
#     after the command.
# 
#     Returns: 
#         tuple of (local_args, s3_downloads, s3_uploads)
# 
#         where new_args contains the new argument list with local paths, 
#         and s3_outputs is a dict mapping local filepaths to s3 paths to 
#         transfer after execution.
#     """
#     s3_uploads = {}
#     s3_downloads = {}
#     local_args = []
#     for arg in args:
#         if not arg.startswith('s3://'):
#             local_args.append(arg)
#             continue
#         
#         src_bucket, src_key = path_to_bucket_and_key(arg)
#         local_tmp = create_tmp_from_key(src_key)
#         
#         # If key exists, add path to downloads, otherwise add path to 
#         # uploads and remove the file or folder
#         if key_exists(src_bucket, src_key):
#             s3_downloads[arg] = local_tmp
#         else:
#             s3_uploads[arg] = local_tmp
#             if rm_local_outpath:
#                 remove_file_or_folder(local_tmp)
# 
#         local_args.append(local_tmp)
# 
#     return local_args, s3_downloads, s3_uploads


# Path: test/test_s3.py
import os
import os.path as op
import shutil
import boto3
import pytest
from subprocess import check_output
from tempfile import mkstemp, mkdtemp
from ob_pipelines.s3 import (
    s3args, key_exists, path_to_bucket_and_key, create_tmp_from_key,
    upload_folder, download_folder, swap_args
)



TEST_BUCKET = os.environ.get('TEST_BUCKET')
OUT_KEY = 'test_out'
IN_KEY = 'test_in'
IN_PREFIX = 'test_prefix/'
OUT_PREFIX = 'test_out_prefix/'
s3 = boto3.client('s3')


@pytest.fixture
def s3_outpath():
    s3.delete_object(Bucket=TEST_BUCKET, Key=OUT_KEY)
    return 's3://{}/{}'.format(TEST_BUCKET, OUT_KEY)



@pytest.fixture