Refactor function selection and test

Author: chahak13

lib/matplotlib/collections.py

@@ -545,14 +545,11 @@ def set_offsets(self, offsets):
         offsets = np.asanyarray(offsets)
         if offsets.shape == (2,):  # Broadcast (2,) -> (1, 2) but nothing else.
             offsets = offsets[None, :]
-        if isinstance(offsets, np.ma.MaskedArray):
-            self._offsets = np.ma.column_stack(
-                (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
-                 np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
-        else:
-            self._offsets = np.column_stack(
-                (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
-                 np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
+        cstack = (np.ma.column_stack if isinstance(offsets, np.ma.MaskedArray)
+                  else np.column_stack)
+        self._offsets = cstack(
+            (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
+                np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
         self.stale = True
 
     def get_offsets(self):

lib/matplotlib/tests/test_collections.py

@@ -1158,15 +1158,14 @@ def test_masked_set_offsets(fig_ref, fig_test):
 
     ax_test = fig_test.add_subplot()
     scat = ax_test.scatter(x, y)
-    x += 1
     scat.set_offsets(np.ma.column_stack([x, y]))
     ax_test.set_xticks([])
     ax_test.set_yticks([])
     ax_test.set_xlim(0, 7)
     ax_test.set_ylim(0, 6)
 
     ax_ref = fig_ref.add_subplot()
-    ax_ref.scatter([2, 3, 6], [1, 2, 5])
+    ax_ref.scatter([1, 2, 5], [1, 2, 5])
     ax_ref.set_xticks([])
     ax_ref.set_yticks([])
     ax_ref.set_xlim(0, 7)
@@ -1180,7 +1179,6 @@ def test_check_offsets_dtype():
 
     fig, ax = plt.subplots()
     scat = ax.scatter(x, y)
-    x += 1
     masked_offsets = np.ma.column_stack([x, y])
     scat.set_offsets(masked_offsets)
     assert isinstance(scat.get_offsets(), type(masked_offsets))