2
\$\begingroup\$

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))
\$\endgroup\$
2
\$\begingroup\$

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")
\$\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.