Include scatter plots in Qt figure options editor.

anntzer · anntzer · commit 731ad32b557a · 2018-11-08T21:19:56.000+01:00
Essentially all the image-handling code can be reused; the only
difference is that collections don't have an `interpolation` field so we
need to check for that.

The rest of the PR is just replacing "image" by "sm" ("ScalarMappable")
throughout...
diff --git a/lib/matplotlib/backends/qt_editor/figureoptions.py b/lib/matplotlib/backends/qt_editor/figureoptions.py
@@ -137,39 +137,43 @@ def prepare_data(d, init):
     # Is there a curve displayed?
     has_curve = bool(curves)
 
-    # Get / Images
-    imagedict = {}
-    for image in axes.get_images():
-        label = image.get_label()
+    # Get ScalarMappables.
+    smdict = {}
+    for sm in [*axes.images, *axes.collections]:
+        label = sm.get_label()
         if label == '_nolegend_':
             continue
-        imagedict[label] = image
-    imagelabels = sorted(imagedict, key=cmp_key)
-    images = []
+        smdict[label] = sm
+    smlabels = sorted(smdict, key=cmp_key)
+    sms = []
     cmaps = [(cmap, name) for name, cmap in sorted(cm.cmap_d.items())]
-    for label in imagelabels:
-        image = imagedict[label]
-        cmap = image.get_cmap()
+    for label in smlabels:
+        sm = smdict[label]
+        cmap = sm.get_cmap()
         if cmap not in cm.cmap_d.values():
-            cmaps = [(cmap, cmap.name)] + cmaps
-        low, high = image.get_clim()
-        imagedata = [
+            cmaps = [(cmap, cmap.name), *cmaps]
+        low, high = sm.get_clim()
+        smdata = [
             ('Label', label),
             ('Colormap', [cmap.name] + cmaps),
             ('Min. value', low),
             ('Max. value', high),
-            ('Interpolation',
-             [image.get_interpolation()]
-             + [(name, name) for name in sorted(mimage.interpolations_names)])]
-        images.append([imagedata, label, ""])
-    # Is there an image displayed?
-    has_image = bool(images)
+        ]
+        if hasattr(sm, "get_interpolation"):  # Images.
+            smdata.append((
+                'Interpolation',
+                [image.get_interpolation(),
+                 *((name, name)
+                   for name in sorted(mimage.interpolations_names))]))
+        sms.append([smdata, label, ""])
+    # Is there a scalarmappable displayed?
+    has_sm = bool(sms)
 
     datalist = [(general, "Axes", "")]
     if curves:
         datalist.append((curves, "Curves", ""))
-    if images:
-        datalist.append((images, "Images", ""))
+    if sms:
+        datalist.append((sms, "Images, etc.", ""))
 
     def apply_callback(data):
         """This function will be called to apply changes"""
@@ -178,7 +182,7 @@ def apply_callback(data):
 
         general = data.pop(0)
         curves = data.pop(0) if has_curve else []
-        images = data.pop(0) if has_image else []
+        sms = data.pop(0) if has_sm else []
         if data:
             raise ValueError("Unexpected field")
 
@@ -223,14 +227,20 @@ def apply_callback(data):
                 line.set_markerfacecolor(markerfacecolor)
                 line.set_markeredgecolor(markeredgecolor)
 
-        # Set / Images
-        for index, image_settings in enumerate(images):
-            image = imagedict[imagelabels[index]]
-            label, cmap, low, high, interpolation = image_settings
-            image.set_label(label)
-            image.set_cmap(cm.get_cmap(cmap))
-            image.set_clim(*sorted([low, high]))
-            image.set_interpolation(interpolation)
+        # Set ScalarMappables.
+        for index, sm_settings in enumerate(sms):
+            sm = smdict[smlabels[index]]
+            sentinel = object()
+            if len(sm_settings) == 5:
+                label, cmap, low, high, interpolation = sm_settings
+            elif len(sm_settings) == 4:
+                label, cmap, low, high = sm_settings
+                interpolation = sentinel
+            sm.set_label(label)
+            sm.set_cmap(cm.get_cmap(cmap))
+            sm.set_clim(*sorted([low, high]))
+            if interpolation is not sentinel:
+                sm.set_interpolation(interpolation)
 
         # re-generate legend, if checkbox is checked
         if generate_legend:
