Fix reading/writing from urllib.request objects

matplotlib · Jan 25, 2016 · 9758da5 · 9758da5
1 parent 25590c8
commit 9758da5
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 10 deletions.
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -517,6 +517,13 @@ def test_minimized_rasterized():
                 assert False
 
 
+@cleanup
+def test_load_from_url():
+    req = six.moves.urllib.request.urlopen(
+        "http://matplotlib.org/_static/logo_sidebar_horiz.png")
+    Z = plt.imread(req)
+
+
 if __name__=='__main__':
     import nose
     nose.runmodule(argv=['-s','--with-doctest'], exit=False)
diff --git a/src/_png.cpp b/src/_png.cpp
@@ -206,7 +206,17 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
             goto exit;
         }
         buff.cursor = 0;
-    } else if ((fp = mpl_PyFile_Dup(py_file, (char *)"wb", &offset))) {
+    } else {
+        #if PY3K
+        if (close_file) {
+        #else
+        if (close_file || PyFile_Check(py_file)) {
+        #endif
+            fp = mpl_PyFile_Dup(py_file, (char *)"wb", &offset);
+        }
+    }
+
+    if (fp) {
         close_dup_file = true;
     } else {
         PyErr_Clear();
@@ -374,10 +384,23 @@ static PyObject *_read_png(PyObject *filein, bool float_result)
         py_file = filein;
     }
 
-    if ((fp = mpl_PyFile_Dup(py_file, (char *)"rb", &offset))) {
+    #if PY3K
+    if (close_file) {
+    #else
+    if (close_file || PyFile_Check(py_file)) {
+    #endif
+        fp = mpl_PyFile_Dup(py_file, (char *)"rb", &offset);
+    }
+
+    if (fp) {
         close_dup_file = true;
+        if (fread(header, 1, 8, fp) != 8) {
+            PyErr_SetString(PyExc_IOError, "error reading PNG header");
+            goto exit;
+        }
     } else {
         PyErr_Clear();
+
         PyObject *read_method = PyObject_GetAttrString(py_file, "read");
         if (!(read_method && PyCallable_Check(read_method))) {
             Py_XDECREF(read_method);
@@ -387,14 +410,6 @@ static PyObject *_read_png(PyObject *filein, bool float_result)
             goto exit;
         }
         Py_XDECREF(read_method);
-    }
-
-    if (fp) {
-        if (fread(header, 1, 8, fp) != 8) {
-            PyErr_SetString(PyExc_IOError, "error reading PNG header");
-            goto exit;
-        }
-    } else {
         _read_png_data(py_file, header, 8);
     }