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.

1 files changed, 35 insertions, 0 deletions

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)