Yes, you could reduce duplication (and make it easier to add new exporters) if you used a dictionary to map extensions to functions. However, you couldn't do it quite as you suggest; you would store the result of the method call in the dictionary, not the method itself. Here is one suggestion, using os.path.splitext
to get the extension defined in save_filename
and functools.partial
to supply the additional argument required for CSVs.
from functools import partial
import os
EXPORTERS = {'.xls': pd.DataFrame.to_excel,
'.xlsx': pd.DataFrame.to_excel,
'.csv': partial(pd.DataFrame.to_csv, index=False)}
def save_to_file(data, save_filename, columns_order=None):
""" Saves data to file
:param data: data to be saved
:param save_filename: the filename
:param columns_order: the order of columns
if needed to rearrange
:return: True if done or None if error
"""
try:
df = pd.DataFrame(data)
if columns_order is not None: # Rearrange columns if needed
df = df[columns_order]
_, ext = os.path.splitext(save_filename)
if ext in EXPORTERS:
EXPORTERS[ext](df, save_filename)
return True
else:
print "Not a supported filetype: {}".format(ext)
except Exception, e:
print "Save_to_file error:", e
Note the following:
- I have stored the methods on the class in the dictionary, allowing it to be defined outside the function, so you need to pass the instance
df
as well as the first argument save_filename
;
- Rather than force the user to specify
""
if the columns_order
isn't required, I have given it a default None
value; and
return None
is implicit.
Including the list of supported filetypes in the docstring and output message is duplication of the information now stored in the dictionary's keys - if you add another supported type, you now need to change three parts of your code.
Finally, you should review your error handling. For one thing, except Exception as e:
is more modern than except Exception, e:
(which is retained on;ly for backwards compatibility, and doesn't work at all in 3.x). Secondly, you should use more specific error handling - what could go wrong, and how should it be dealt with? At the very least, I would differentiate between:
- The conversion to
DataFrame
or column rearrangement not working; and
- A failure in file operations.
Thirdly, rather than return
either True
or None
, it would be conventional to raise an error from your function in the cases where you currently print
a message - if the data hasn't been saved, the calling function should do something about it - and return None
(again, implicitly) otherwise.