1
\$\begingroup\$

I'm trying to copy files to S3 after writing them locally: I have 2 functions which I need to do this. What is the best way to refactor the code that exists in both.

    s3 = False
    if data_fp.lower().startswith('s3://'):
        s3_data_fp = data_fp
        data_fp = os.path.join(_TMP, os.path.basename(data_fp))
        s3 = True

def save_json(data_fp, data, sort_keys=True, indent=4):
    s3 = False
    if data_fp.lower().startswith('s3://'):
        s3_data_fp = data_fp
        data_fp = os.path.join(_TMP, os.path.basename(data_fp))
        s3 = True

    with open(data_fp, 'w') as output_file:
        json.dump(data, output_file, cls=NumpyEncoder, sort_keys=sort_keys,
                  indent=indent)
    if s3:
        save_in_s3(data_fp, s3_data_fp)


# def save_hdf5(data_fp: str, data: Dict[str, object]):
def save_hdf5(data_fp, data, metadata=None):
    s3 = False
    if data_fp.lower().startswith('s3://'):
        s3_data_fp = data_fp
        data_fp = os.path.join(_TMP, os.path.basename(data_fp))
        s3 = True

    if metadata is None:
        metadata = {}
    mode = 'w'
    if os.path.isfile(data_fp):
        mode = 'r+'
    with h5py.File(data_fp, mode) as h5_file:
        for key, value in data.items():
            dataset = h5_file.create_dataset(key, data=value)                

    if s3:
        save_in_s3(data_fp, s3_data_fp)
\$\endgroup\$
3
\$\begingroup\$

The S3-functionality wraps the specific code, which does not change the variables. So how about a decorator?

from functools import wraps

def handle_s3_upload(fcn):
    @wraps(fcn)
    def inner(data_fp, data, *args, **kwargs):
        # handle S3 path
        s3 = False
        if data_fp.lower().startswith('s3://'):
            s3_data_fp = data_fp
            data_fp = os.path.join(_TMP, os.path.basename(data_fp))
            s3 = True
        # call decorated function
        fcn(data_fp, data, *args, **kwargs)
        # S3 upload
        if s3:
            save_in_s3(data_fp, s3_data_fp)
    return inner

@handle_s3_upload
def save_json(data_fp, data, sort_keys=True, indent=4):
    with open(data_fp, 'w') as output_file:
        json.dump(data, output_file, cls=NumpyEncoder, sort_keys=sort_keys,
                  indent=indent)

@handle_s3_upload
def save_hdf5(data_fp, data, metadata=None):
    if metadata is None:
        metadata = {}
    mode = 'w'
    if os.path.isfile(data_fp):
        mode = 'r+'
    with h5py.File(data_fp, mode) as h5_file:
        for key, value in data.items():
            dataset = h5_file.create_dataset(key, data=value)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.