ENH: Add provision for headers and footers to savetxt, fixes ticket 1236.

I suggest using a separate keyword argument for structured arrays. It might
also be nice to be able to add a manual header.

2 files changed, 57 insertions, 4 deletions

diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py
index 6c4a9942f..30c8cc5be 100644
--- a/numpy/lib/npyio.py
+++ b/numpy/lib/npyio.py
@@ -834,7 +834,8 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None,
         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.
 
@@ -846,14 +847,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.
 
@@ -977,8 +989,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 8bfb64cee..54c2ddabb 100644
--- a/numpy/lib/tests/test_io.py
+++ b/numpy/lib/tests/test_io.py
@@ -227,6 +227,41 @@ class TestSaveTxt(TestCase):
         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)