diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py index 44a91c0c7995..61fdf25587f7 100644 --- a/numpy/lib/npyio.py +++ b/numpy/lib/npyio.py @@ -833,7 +833,8 @@ def split_line(line): return X -def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n'): +def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', + footer='', comments='# '): """ Save an array to a text file. @@ -845,14 +846,25 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n'): transparently. X : array_like Data to be saved to a text file. - fmt : str or sequence of strs + fmt : str or sequence of strs, optional A single format (%10.5f), a sequence of formats, or a multi-format string, e.g. 'Iteration %d -- %10.5f', in which case `delimiter` is ignored. - delimiter : str + delimiter : str, optional Character separating columns. - newline : str + newline : str, optional .. versionadded:: 1.5.0 + header : str, optional + String that will be written at the beginning of the file. + .. versionadded:: 2.0.0 + footer : str, optional + String that will be written at the end of the file. + .. versionadded:: 2.0.0 + comments : str, optional + String that will be prepended to the ``header`` and ``footer`` strings, + to mark them as comments. Default: '# ', as expected by e.g. + ``numpy.loadtxt``. + .. versionadded:: 2.0.0 Character separating lines. @@ -976,8 +988,14 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n'): else: format = fmt + if len(header) > 0: + header = header.replace('\n', '\n' + comments) + fh.write(asbytes(comments + header + newline)) for row in X: fh.write(asbytes(format % tuple(row) + newline)) + if len(footer) > 0: + footer = footer.replace('\n', '\n' + comments) + fh.write(asbytes(comments + footer + newline)) finally: if own_fh: fh.close() diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py index 18585375e387..e108d538011a 100644 --- a/numpy/lib/tests/test_io.py +++ b/numpy/lib/tests/test_io.py @@ -227,6 +227,41 @@ def test_format(self): lines = c.readlines() assert_equal(lines, asbytes_nested(['01 : 2.0\n', '03 : 4.0\n'])) + def test_header_footer(self): + """ + Test the functionality of the header and footer keyword argument. + """ + c = StringIO() + a = np.array([(1, 2), (3, 4)], dtype=np.int) + test_header_footer = 'Test header / footer' + # Test the header keyword argument + np.savetxt(c, a, fmt='%1d', header=test_header_footer) + c.seek(0) + assert_equal(c.read(), + asbytes('# ' + test_header_footer +'\n1 2\n3 4\n' )) + # Test the footer keyword argument + c = StringIO() + np.savetxt(c, a, fmt='%1d', footer=test_header_footer) + c.seek(0) + assert_equal(c.read(), + asbytes('1 2\n3 4\n# ' + test_header_footer + '\n')) + # Test the commentstr keyword argument used on the header + c = StringIO() + commentstr = '% ' + np.savetxt(c, a, fmt='%1d', header=test_header_footer, + comments=commentstr) + c.seek(0) + assert_equal(c.read(), + asbytes(commentstr + test_header_footer + '\n' + '1 2\n3 4\n')) + # Test the commentstr keyword argument used on the footer + c = StringIO() + commentstr = '% ' + np.savetxt(c, a, fmt='%1d', footer=test_header_footer, + comments=commentstr) + c.seek(0) + assert_equal(c.read(), + asbytes('1 2\n3 4\n' + commentstr + test_header_footer + '\n')) + def test_file_roundtrip(self): f, name = mkstemp() os.close(f)