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")