Merge pull request #6222 from madphysicist/setp-output

tacaswell · tacaswell · commit c06bb2b36c39 · 2016-03-26T09:35:30.000-04:00
ENH: Added `file` keyword to `setp` to redirect output
diff --git a/doc/users/whats_new/setp_output.rst b/doc/users/whats_new/setp_output.rst
@@ -0,0 +1,7 @@
+`Artist.setp` (and `pyplot.setp`) accept a `file` argument
+----------------------------------------------------------
+
+The argument is keyword-only. It allows an output file other than
+`sys.stdout` to be specified. It works exactly like the `file` argument
+to `print`.
+
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -1521,6 +1521,12 @@ def setp(obj, *args, **kwargs):
       >>> setp(line)
           ... long output listing omitted
 
+    You may specify another output file to `setp` if `sys.stdout` is not
+    acceptable for some reason using the `file` keyword-only argument::
+
+      >>> with fopen('output.log') as f:
+      >>>     setp(line, file=f)
+
     :func:`setp` operates on a single instance or a iterable of
     instances. If you are in query mode introspecting the possible
     values, only the first instance in the sequence is used. When
@@ -1548,12 +1554,16 @@ def setp(obj, *args, **kwargs):
 
     insp = ArtistInspector(objs[0])
 
-    if len(kwargs) == 0 and len(args) == 0:
-        print('\n'.join(insp.pprint_setters()))
-        return
+    # file has to be popped before checking if kwargs is empty
+    printArgs = {}
+    if 'file' in kwargs:
+        printArgs['file'] = kwargs.pop('file')
 
-    if len(kwargs) == 0 and len(args) == 1:
-        print(insp.pprint_setters(prop=args[0]))
+    if not kwargs and len(args) < 2:
+        if args:
+            print(insp.pprint_setters(prop=args[0]), **printArgs)
+        else:
+            print('\n'.join(insp.pprint_setters()), **printArgs)
         return
 
     if len(args) % 2:
diff --git a/lib/matplotlib/tests/test_artist.py b/lib/matplotlib/tests/test_artist.py
@@ -211,12 +211,18 @@ def test_properties():
 
 @cleanup
 def test_setp():
+    # Check arbitrary iterables
     fig, axes = plt.subplots()
     lines1 = axes.plot(range(3))
     lines2 = axes.plot(range(3))
     martist.setp(chain(lines1, lines2), 'lw', 5)
     plt.setp(axes.spines.values(), color='green')
 
+    # Check `file` argument
+    sio = io.StringIO()
+    plt.setp(lines1, 'zorder', file=sio)
+    assert sio.getvalue() == '  zorder: any number \n'
+
 
 if __name__ == '__main__':
     import nose
