Merge pull request #17282 from anntzer/segment_hits_overflow

timhoffm · web-flow · commit 480ea3f5da48 · 2020-05-01T16:39:18.000+02:00
Don't divide by zero in Line2D.segment_hits.
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -80,7 +80,7 @@ def segment_hits(cx, cy, x, y, radius):
     """
     # Process single points specially
     if len(x) <= 1:
-        res, = np.nonzero((cx - x) ** 2 + (cy - y) ** 2 <= radius ** 2)
+        res, = np.nonzero(np.hypot(cx - x, cy - y) <= radius)
         return res
 
     # We need to lop the last element off a lot.
@@ -89,24 +89,24 @@ def segment_hits(cx, cy, x, y, radius):
     # Only look at line segments whose nearest point to C on the line
     # lies within the segment.
     dx, dy = x[1:] - xr, y[1:] - yr
-    Lnorm_sq = dx ** 2 + dy ** 2  # Possibly want to eliminate Lnorm==0
-    u = ((cx - xr) * dx + (cy - yr) * dy) / Lnorm_sq
-    candidates = (u >= 0) & (u <= 1)
+    u = (cx - xr) * dx + (cy - yr) * dy
+    candidates = (0 <= u) & (u <= dx ** 2 + dy ** 2)
 
     # Note that there is a little area near one side of each point
     # which will be near neither segment, and another which will
     # be near both, depending on the angle of the lines.  The
     # following radius test eliminates these ambiguities.
-    point_hits = (cx - x) ** 2 + (cy - y) ** 2 <= radius ** 2
+    point_hits = np.hypot(cx - x, cy - y) <= radius
     candidates = candidates & ~(point_hits[:-1] | point_hits[1:])
 
     # For those candidates which remain, determine how far they lie away
     # from the line.
     px, py = xr + u * dx, yr + u * dy
-    line_hits = (cx - px) ** 2 + (cy - py) ** 2 <= radius ** 2
+    line_hits = np.hypot(cx - px, cy - py) <= radius
     line_hits = line_hits & candidates
-    points, = point_hits.ravel().nonzero()
-    lines, = line_hits.ravel().nonzero()
+
+    points, = point_hits.nonzero()
+    lines, = line_hits.nonzero()
     return np.concatenate((points, lines))
 
 
@@ -454,21 +454,16 @@ def contains(self, mouseevent):
         else:
             pixels = self.figure.dpi / 72. * self.pickradius
 
-        # The math involved in checking for containment (here and inside of
-        # segment_hits) assumes that it is OK to overflow, so temporarily set
-        # the error flags accordingly.
-        with np.errstate(all='ignore'):
-            # Check for collision
-            if self._linestyle in ['None', None]:
-                # If no line, return the nearby point(s)
-                ind, = np.nonzero(
-                    (xt - mouseevent.x) ** 2 + (yt - mouseevent.y) ** 2
-                    <= pixels ** 2)
-            else:
-                # If line, return the nearby segment(s)
-                ind = segment_hits(mouseevent.x, mouseevent.y, xt, yt, pixels)
-                if self._drawstyle.startswith("steps"):
-                    ind //= 2
+        # Check for collision
+        if self._linestyle in ['None', None]:
+            # If no line, return the nearby point(s)
+            ind, = np.nonzero(
+                np.hypot(xt - mouseevent.x, yt - mouseevent.y) <= pixels)
+        else:
+            # If line, return the nearby segment(s)
+            ind = segment_hits(mouseevent.x, mouseevent.y, xt, yt, pixels)
+            if self._drawstyle.startswith("steps"):
+                ind //= 2
 
         ind += self.ind_offset
 
