From 9e7483ffc3d70891bee6fecdbeaa097658835a34 Mon Sep 17 00:00:00 2001
From: Zac-HD <zac.hatfield.dodds@gmail.com>
Date: Mon, 6 Sep 2021 23:36:33 +1000
Subject: [PATCH] Trim more tracebacks

---
 hypothesis-python/src/hypothesis/errors.py     | 18 +++++++++++-------
 .../src/hypothesis/internal/escalation.py      |  7 +++----
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/errors.py b/hypothesis-python/src/hypothesis/errors.py
index c342d8ea5e..674089a211 100644
--- a/hypothesis-python/src/hypothesis/errors.py
+++ b/hypothesis-python/src/hypothesis/errors.py
@@ -18,6 +18,10 @@ class HypothesisException(Exception):
     """Generic parent class for exceptions thrown by Hypothesis."""
 
 
+class _Trimmable(HypothesisException):
+    """Hypothesis can trim these tracebacks even if they're raised internally."""
+
+
 class CleanupFailed(HypothesisException):
     """At least one cleanup task failed and no other exception was raised."""
 
@@ -40,7 +44,7 @@ def __init__(self, condition_string, extra=""):
         super().__init__(f"No examples found of condition {condition_string}{extra}")
 
 
-class Unsatisfiable(HypothesisException):
+class Unsatisfiable(_Trimmable):
     """We ran out of time or examples before we could find enough examples
     which satisfy the assumptions of this hypothesis.
 
@@ -53,7 +57,7 @@ class Unsatisfiable(HypothesisException):
     """
 
 
-class Flaky(HypothesisException):
+class Flaky(_Trimmable):
     """This function appears to fail non-deterministically: We have seen it
     fail when passed this example at least once, but a subsequent invocation
     did not fail.
@@ -70,7 +74,7 @@ class Flaky(HypothesisException):
     """
 
 
-class InvalidArgument(HypothesisException, TypeError):
+class InvalidArgument(_Trimmable, TypeError):
     """Used to indicate that the arguments to a Hypothesis function were in
     some manner incorrect."""
 
@@ -88,7 +92,7 @@ class InvalidState(HypothesisException):
     """The system is not in a state where you were allowed to do that."""
 
 
-class InvalidDefinition(HypothesisException, TypeError):
+class InvalidDefinition(_Trimmable, TypeError):
     """Used to indicate that a class definition was not well put together and
     has something wrong with it."""
 
@@ -97,7 +101,7 @@ class HypothesisWarning(HypothesisException, Warning):
     """A generic warning issued by Hypothesis."""
 
 
-class FailedHealthCheck(HypothesisWarning):
+class FailedHealthCheck(_Trimmable):
     """Raised when a test fails a preliminary healthcheck that occurs before
     execution."""
 
@@ -129,12 +133,12 @@ class Frozen(HypothesisException):
     after freeze() has been called."""
 
 
-class MultipleFailures(HypothesisException):
+class MultipleFailures(_Trimmable):
     """Indicates that Hypothesis found more than one distinct bug when testing
     your code."""
 
 
-class DeadlineExceeded(HypothesisException):
+class DeadlineExceeded(_Trimmable):
     """Raised when an individual test body has taken too long to run."""
 
     def __init__(self, runtime, deadline):
diff --git a/hypothesis-python/src/hypothesis/internal/escalation.py b/hypothesis-python/src/hypothesis/internal/escalation.py
index 946178fae2..a9860f94cf 100644
--- a/hypothesis-python/src/hypothesis/internal/escalation.py
+++ b/hypothesis-python/src/hypothesis/internal/escalation.py
@@ -23,11 +23,10 @@
 import hypothesis
 from hypothesis.errors import (
     DeadlineExceeded,
-    Flaky,
     HypothesisException,
-    MultipleFailures,
     StopTest,
     UnsatisfiedAssumption,
+    _Trimmable,
 )
 
 
@@ -92,10 +91,10 @@ def get_trimmed_traceback(exception=None):
     # was raised inside Hypothesis (and is not a MultipleFailures)
     if hypothesis.settings.default.verbosity >= hypothesis.Verbosity.debug or (
         is_hypothesis_file(traceback.extract_tb(tb)[-1][0])
-        and not isinstance(exception, (Flaky, MultipleFailures))
+        and not isinstance(exception, _Trimmable)
     ):
         return tb
-    while tb is not None and (
+    while tb.tb_next is not None and (
         # If the frame is from one of our files, it's been added by Hypothesis.
         is_hypothesis_file(getframeinfo(tb.tb_frame)[0])
         # But our `@proxies` decorator overrides the source location,