Convenience function for writing files in CSV, JSON, etc

Question

The McCabe code complexity of the following function is 19 as found by pytest --mccabe test.py (from my toy-project mpu). While I agree that this function is length and likely will become longer in future, I don't know how to improve the situation.

What it does

The code mpu.io.write(filepath) is a convenience function for writing a few common file formats.

Code

This code should work in Python 2 and 3.

def write(filepath, data, **kwargs):
    """
    Write a file.
    Supported formats:
    * CSV
    * JSON, JSONL
    * pickle
    Parameters
    ----------
    filepath : str
        Path to the file that should be read. This methods action depends
        mainly on the file extension.
    data : dict or list
        Content that should be written
    kwargs : dict
        Any keywords for the specific file format.
    Returns
    -------
    data : str or bytes
    """
    if filepath.lower().endswith('.csv'):
        kwargs_open = {'newline': ''}
        mode = 'w'
        if sys.version_info < (3, 0):
            kwargs_open.pop('newline', None)
            mode = 'wb'
        with open(filepath, mode, **kwargs_open) as fp:
            if 'delimiter' not in kwargs:
                kwargs['delimiter'] = ','
            if 'quotechar' not in kwargs:
                kwargs['quotechar'] = '"'
            with open(filepath, 'w') as fp:
                writer = csv.writer(fp, **kwargs)
                writer.writerows(data)
        return data
    elif filepath.lower().endswith('.json'):
        with io_stl.open(filepath, 'w', encoding='utf8') as outfile:
            if 'indent' not in kwargs:
                kwargs['indent'] = 4
            if 'sort_keys' not in kwargs:
                kwargs['sort_keys'] = True
            if 'separators' not in kwargs:
                kwargs['separators'] = (',', ': ')
            if 'ensure_ascii' not in kwargs:
                kwargs['ensure_ascii'] = False
            str_ = json.dumps(data, **kwargs)
            outfile.write(to_unicode(str_))
    elif filepath.lower().endswith('.jsonl'):
        print(filepath)
        with io_stl.open(filepath, 'w', encoding='utf8') as outfile:
            kwargs['indent'] = None  # JSON has to be on one line!
            if 'sort_keys' not in kwargs:
                kwargs['sort_keys'] = True
            if 'separators' not in kwargs:
                kwargs['separators'] = (',', ': ')
            if 'ensure_ascii' not in kwargs:
                kwargs['ensure_ascii'] = False
            for line in data:
                str_ = json.dumps(line, **kwargs)
                outfile.write(to_unicode(str_))
                outfile.write(u'\n')
    elif filepath.lower().endswith('.pickle'):
        if 'protocol' not in kwargs:
            kwargs['protocol'] = pickle.HIGHEST_PROTOCOL
        with open(filepath, 'wb') as handle:
            pickle.dump(data, handle, **kwargs)
    elif (filepath.lower().endswith('.yml') or
          filepath.lower().endswith('.yaml')):
        raise NotImplementedError('YAML is not supported, because you need '
                                  'PyYAML in Python3. '
                                  'See '
                                  'https://stackoverflow.com/a/42054860/562769'
                                  ' as a guide how to use it.')
    elif (filepath.lower().endswith('.h5') or
          filepath.lower().endswith('.hdf5')):
        raise NotImplementedError('HDF5 is not supported. See '
                                  'https://stackoverflow.com/a/41586571/562769'
                                  ' as a guide how to use it.')
    else:
        raise NotImplementedError('File \'{}\' is not known.'.format(filepath))

Answer 1

You’ve got your code split into various sections; why not make the sections their own methods?

if filepath.lower().endswith('.csv'):
   return _write_csv(filepath, data, kwargs)
elif filepath.lower().endswith('.json'):
   return _write_json(filepath, data, kwargs)
elif filepath.lower().endswith('.jsonl'):
   return _write_jsonl(filepath, data, kwargs)
elif filepath.lower().endswith('.pickle'):
   return _write_pickle(filepath, data, kwargs)
elif filepath.lower().endswith('.yml') or
     filepath.lower().endswith('.yaml'):
   return _write_yml(filepath, data, kwargs)
elif filepath.lower().endswith('.h5') or
     filepath.lower().endswith('.hdf5'):
   return _write_hdf5(filepath, data, kwargs)
else:
   raise NotImplementedError('File \'{}\' is not known.'.format(filepath))

With the appropriate dict, you could even:

path = filepath.lower()
for ending, writer in writers.items():
    if path.endswith(ending):
        return writer(filepath, data, kwargs)
raise NotImplementedError(f"File {filepath} is not known")

Or,

ext = os.path.splitext(filepath)[1].lower()
if ext in writers:
    writers[ext](filepath, data, kwargs)
raise NotImplementedError(f"File {filepath} is not known")