Improve debugging for the SB() context manager

Author: mdmintz

seleniumbase/core/colored_traceback.py

@@ -0,0 +1,67 @@
+import sys
+
+
+def add_hook(always=False, style="default", debug=False):
+    import os
+
+    if os.environ.get("NO_COLOR", ""):
+        return  # https://no-color.org
+    isatty = getattr(sys.stderr, "isatty", lambda: False)
+    if always or isatty():
+        colorizer = Colorizer(style, debug)
+        sys.excepthook = colorizer.colorize_traceback
+
+
+class Colorizer(object):
+    def __init__(self, style, debug=False):
+        self.style = style
+        self.debug = debug
+
+    def colorize_traceback(self, type, value, tb):
+        import traceback
+        import pygments.lexers
+
+        tb_text = "".join(traceback.format_exception(type, value, tb))
+        lexer_name = "pytb" if sys.version_info < (3, ) else "py3tb"
+        lexer = pygments.lexers.get_lexer_by_name(lexer_name)
+        tb_colored = pygments.highlight(tb_text, lexer, self.formatter)
+        self.stream.write(tb_colored)
+
+    @property
+    def formatter(self):
+        from pygments.formatters import get_formatter_by_name
+        import pygments.util
+
+        colors = _get_term_color_support()
+        if self.debug:
+            sys.stderr.write("Detected support for %s colors\n" % colors)
+        if colors == 256:
+            fmt_options = {"style": self.style}
+        elif self.style in ("light", "dark"):
+            fmt_options = {"bg": self.style}
+        else:
+            fmt_options = {"bg": "dark"}
+        fmt_alias = "terminal256" if colors == 256 else "terminal"
+        try:
+            return get_formatter_by_name(fmt_alias, **fmt_options)
+        except pygments.util.ClassNotFound as ex:
+            if self.debug:
+                sys.stderr.write(str(ex) + "\n")
+            return get_formatter_by_name(fmt_alias)
+
+    @property
+    def stream(self):
+        try:
+            import colorama
+        except ImportError:
+            return sys.stderr
+        return colorama.AnsiToWin32(sys.stderr)
+
+
+def _get_term_color_support():
+    try:
+        import curses
+    except ImportError:
+        return 16
+    curses.setupterm()
+    return curses.tigetnum("colors")

seleniumbase/plugins/sb_manager.py

@@ -6,7 +6,7 @@
 Usage --> ``with SB() as sb:``
 Usage example -->
     from seleniumbase import SB
-    with SB() as sb:  # Lots of args! (Eg. headless=True)
+    with SB() as sb:  # Many args! Eg. SB(browser="edge")
         sb.open("https://google.com/ncr")
         sb.type('[name="q"]', "SeleniumBase on GitHub\n")
         sb.click('a[href*="github.com/seleniumbase"]')
@@ -109,10 +109,12 @@ def SB(
     from seleniumbase import BaseCase
     from seleniumbase import config as sb_config
     from seleniumbase.config import settings
+    from seleniumbase.core import colored_traceback
     from seleniumbase.fixtures import constants
     from seleniumbase.fixtures import shared_utils
 
     sb_config_backup = sb_config
+    sb_config._do_sb_post_mortem = False
     sys_argv = sys.argv
     archive_logs = False
     existing_runner = False
@@ -753,6 +755,7 @@ def SB(
         sb.headless_active = False
     test_name = None
     terminal_width = shared_utils.get_terminal_width()
+    colored_traceback.add_hook()
     if test:
         import colorama
         import os
@@ -787,6 +790,11 @@ def SB(
     sb.setUp()
     test_passed = True  # This can change later
     teardown_exception = None
+    if "--trace" in sys_argv:
+        import pdb
+
+        pdb.set_trace()  # Debug Mode
+        # Type "s" and press [Enter] to step into "yield sb".
     try:
         yield sb
     except Exception as e:
@@ -803,8 +811,10 @@ def SB(
                 sb.cm_filename = filename
             except Exception:
                 sb.cm_filename = None
-        # Tests will raise an exception later if "raise_test_failure"
+        # Tests will raise an exception if raise_test_failure is True
     finally:
+        if sb._has_failure and "--pdb" in sys_argv:
+            sb_config._do_sb_post_mortem = True
         try:
             sb.tearDown()
         except Exception as t_e:
@@ -821,7 +831,7 @@ def SB(
             sb_config._context_of_runner = True
         if test_name:
             result = "passed"
-            if not test_passed:
+            if test and not test_passed:
                 result = "failed"
                 c1 = colorama.Fore.RED
             end_text = (
@@ -840,7 +850,7 @@ def SB(
                 left_space = left_spaces * "="
                 right_spaces = remaining_spaces - left_spaces
                 right_space = right_spaces * "="
-            if not test_passed:
+            if test and not test_passed:
                 print(the_traceback)
             if not test_name.startswith("runpy.py:"):
                 print(