Merge pull request #5910 from mdboom/image-read-from-url

WeatherGod · WeatherGod · commit 3675841f686d · 2016-01-25T15:11:01.000-05:00
Fix reading/writing from urllib.request objects
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -5,6 +5,8 @@
 import io
 import os
 
+from nose.plugins.attrib import attr
+
 import numpy as np
 
 from matplotlib.testing.decorators import (image_comparison,
@@ -517,6 +519,13 @@ def test_minimized_rasterized():
                 assert False
 
 
+@attr('network')
+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);
     }
 
