Merge pull request numpy#13409 from seberg/fix-unicode-fmt-savetxt

mattip · web-flow · commit bf1e9b747db5 · 2019-04-26T14:47:48.000-07:00
BUG: (py2 only) fix unicode support for savetxt fmt string
diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py
@@ -1399,7 +1399,7 @@ def first_write(self, v):
             if len(fmt) != ncol:
                 raise AttributeError('fmt has wrong shape.  %s' % str(fmt))
             format = asstr(delimiter).join(map(asstr, fmt))
-        elif isinstance(fmt, str):
+        elif isinstance(fmt, basestring):
             n_fmt_chars = fmt.count('%')
             error = ValueError('fmt has wrong number of %% formats:  %s' % fmt)
             if n_fmt_chars == 1:
diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py
@@ -561,6 +561,19 @@ def test_unicode_stringstream(self):
         s.seek(0)
         assert_equal(s.read(), utf8 + '\n')
 
+    @pytest.mark.parametrize("fmt", [u"%f", b"%f"])
+    @pytest.mark.parametrize("iotype", [StringIO, BytesIO])
+    def test_unicode_and_bytes_fmt(self, fmt, iotype):
+        # string type of fmt should not matter, see also gh-4053
+        a = np.array([1.])
+        s = iotype()
+        np.savetxt(s, a, fmt=fmt)
+        s.seek(0)
+        if iotype is StringIO:
+            assert_equal(s.read(), u"%f\n" % 1.)
+        else:
+            assert_equal(s.read(), b"%f\n" % 1.)
+
 
 class LoadTxtBase(object):
     def check_compressed(self, fopen, suffixes):
