Refactor matplotlib_handler to use direct import for Axes and improve…

… exception handling in stage_test
hyperskill · Feb 15, 2025 · f59edb2 · f59edb2
1 parent 7121c10
commit f59edb2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/hstest/stage/stage_test.py b/hstest/stage/stage_test.py
@@ -174,7 +174,8 @@ def run_tests(self, *, debug=False, is_unittest: bool = False) -> tuple[int, str
             SystemHandler.tear_down()
             return passed(is_unittest)
 
-        except BaseException:
+        except BaseException as caught_ex:
+            ex = caught_ex
             if need_tear_down:
                 try:
                     StageTest.curr_test_run.tear_down()

diff --git a/hstest/testing/plotting/matplotlib_handler.py b/hstest/testing/plotting/matplotlib_handler.py
@@ -199,7 +199,9 @@ def boxplot(x, **kwargs) -> None:
 
         import matplotlib as mpl
 
-        class CustomMatplotlibAxes(mpl.axes.Axes):
+        import matplotlib.axes
+
+        class CustomMatplotlibAxes(matplotlib.axes.Axes):
             def hist(self, x, *a, **kw) -> None:
                 hist(x, *a, **kw)
 
@@ -227,17 +229,15 @@ def imshow(self, x, **kwargs) -> None:
             def boxplot(self, x, **kwargs) -> None:
                 boxplot(x, **kwargs)
 
-        import matplotlib as mpl
-
         if not MatplotlibHandler._saved:
-            MatplotlibHandler._Axes = deepcopy(mpl.axes.Axes)
+            MatplotlibHandler._Axes = deepcopy(matplotlib.axes.Axes)
 
         # should be replaced before import matplotlib.pyplot as plt
-        mpl.axes.Axes = CustomMatplotlibAxes
+        matplotlib.axes.Axes = CustomMatplotlibAxes
 
         from matplotlib.projections import projection_registry
 
-        projection_registry.register(mpl.axes.Axes)
+        projection_registry.register(CustomMatplotlibAxes)
 
         import matplotlib.pyplot as plt