[library/python/pytest] Better py3 colorization

9 files changed, 130 insertions, 12 deletions

diff --git a/library/python/pytest/plugins/ya.py b/library/python/pytest/plugins/ya.py
index 723904df92..e5d3ec09e5 100644
--- a/library/python/pytest/plugins/ya.py
+++ b/library/python/pytest/plugins/ya.py
@@ -617,19 +617,9 @@ def colorize(longrepr):
             return io.getvalue().strip()
         return yatest_lib.tools.to_utf8(longrepr)
 
+    # Use arcadia style colorization
     text = yatest_lib.tools.to_utf8(longrepr)
-    pos = text.find("E   ")
-    if pos == -1:
-        return text
-
-    bt, error = text[:pos], text[pos:]
-    filters = [
-        # File path, line number and function name
-        (re.compile(r"^(.*?):(\d+): in (\S+)", flags=re.MULTILINE), r"[[unimp]]\1[[rst]]:[[alt2]]\2[[rst]]: in [[alt1]]\3[[rst]]"),
-    ]
-    for regex, substitution in filters:
-        bt = regex.sub(substitution, bt)
-    return "{}[[bad]]{}".format(bt, error)
+    return tools.colorize_pytest_error(text)
 
 
 class TestItem(object):
diff --git a/library/python/pytest/ut/canondata/py2test/result.json b/library/python/pytest/ut/canondata/py2test/result.json
new file mode 100644
index 0000000000..521cc45c6d
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py2test/result.json
@@ -0,0 +1,8 @@
+{
+    "test_tools.test_colorize_pytest_error[nested_by]": {
+        "uri": "file://test_tools.test_colorize_pytest_error_nested_by_/extracted"
+    },
+    "test_tools.test_colorize_pytest_error[simple_bt]": {
+        "uri": "file://test_tools.test_colorize_pytest_error_simple_bt_/extracted"
+    }
+}
diff --git a/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_nested_by_/extracted b/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_nested_by_/extracted
new file mode 100644
index 0000000000..08effe9f30
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_nested_by_/extracted
@@ -0,0 +1,15 @@
+
+[[unimp]]path/test.py[[rst]]:[[alt2]]13[[rst]]: in [[alt1]]test[[rst]]
+    raise Exception("123
+123
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test")[[rst]]
+[[bad]]E   Exception: 123
+E   123
+E   path/test.py:15: in test
+
+During handling of the above exception, another exception occurred:
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test[[rst]]
+    return foo(0)
+[[unimp]]path/test.py[[rst]]:[[alt2]]5[[rst]]: in [[alt1]]foo[[rst]]
+    b, a = 1, 1 / p
+[[bad]]E   ZeroDivisionError: division by zero
diff --git a/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_simple_bt_/extracted b/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_simple_bt_/extracted
new file mode 100644
index 0000000000..41d6246500
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py2test/test_tools.test_colorize_pytest_error_simple_bt_/extracted
@@ -0,0 +1,6 @@
+
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test[[rst]]
+    return foo(0)
+[[unimp]]path/test.py[[rst]]:[[alt2]]5[[rst]]: in [[alt1]]foo[[rst]]
+    b, a = 1, 1 / p
+[[bad]]E   ZeroDivisionError: integer division or modulo by zero
diff --git a/library/python/pytest/ut/canondata/py3test/result.json b/library/python/pytest/ut/canondata/py3test/result.json
new file mode 100644
index 0000000000..521cc45c6d
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py3test/result.json
@@ -0,0 +1,8 @@
+{
+    "test_tools.test_colorize_pytest_error[nested_by]": {
+        "uri": "file://test_tools.test_colorize_pytest_error_nested_by_/extracted"
+    },
+    "test_tools.test_colorize_pytest_error[simple_bt]": {
+        "uri": "file://test_tools.test_colorize_pytest_error_simple_bt_/extracted"
+    }
+}
diff --git a/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_nested_by_/extracted b/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_nested_by_/extracted
new file mode 100644
index 0000000000..08effe9f30
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_nested_by_/extracted
@@ -0,0 +1,15 @@
+
+[[unimp]]path/test.py[[rst]]:[[alt2]]13[[rst]]: in [[alt1]]test[[rst]]
+    raise Exception("123
+123
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test")[[rst]]
+[[bad]]E   Exception: 123
+E   123
+E   path/test.py:15: in test
+
+During handling of the above exception, another exception occurred:
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test[[rst]]
+    return foo(0)
+[[unimp]]path/test.py[[rst]]:[[alt2]]5[[rst]]: in [[alt1]]foo[[rst]]
+    b, a = 1, 1 / p
+[[bad]]E   ZeroDivisionError: division by zero
diff --git a/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_simple_bt_/extracted b/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_simple_bt_/extracted
new file mode 100644
index 0000000000..41d6246500
--- /dev/null
+++ b/library/python/pytest/ut/canondata/py3test/test_tools.test_colorize_pytest_error_simple_bt_/extracted
@@ -0,0 +1,6 @@
+
+[[unimp]]path/test.py[[rst]]:[[alt2]]15[[rst]]: in [[alt1]]test[[rst]]
+    return foo(0)
+[[unimp]]path/test.py[[rst]]:[[alt2]]5[[rst]]: in [[alt1]]foo[[rst]]
+    b, a = 1, 1 / p
+[[bad]]E   ZeroDivisionError: integer division or modulo by zero
diff --git a/library/python/pytest/ut/test_tools.py b/library/python/pytest/ut/test_tools.py
index 412ec42d24..212a20fb17 100644
--- a/library/python/pytest/ut/test_tools.py
+++ b/library/python/pytest/ut/test_tools.py
@@ -95,3 +95,39 @@ def test_path_resolving_for_local_conftest_load_policy(
     mocker.patch.object(sys, 'extra_modules', extra_modules)
     got = yatest_tools.split_node_id(node_id + parameters)
     assert (expected_class_name, expected_test_name + parameters) == got
+
+
+DATA = [
+    (
+        "simple_bt",
+        """
+path/test.py:15: in test
+    return foo(0)
+path/test.py:5: in foo
+    b, a = 1, 1 / p
+E   ZeroDivisionError: integer division or modulo by zero
+""",
+    ),
+    (
+        "nested_by",
+        """
+path/test.py:13: in test
+    raise Exception("123\n123\npath/test.py:15: in test")
+E   Exception: 123
+E   123
+E   path/test.py:15: in test
+
+During handling of the above exception, another exception occurred:
+path/test.py:15: in test
+    return foo(0)
+path/test.py:5: in foo
+    b, a = 1, 1 / p
+E   ZeroDivisionError: division by zero
+""",
+    ),
+]
+
+
+@pytest.mark.parametrize("bt", [x[1] for x in DATA], ids=([x[0] for x in DATA]))
+def test_colorize_pytest_error(bt):
+    return yatest_tools.colorize_pytest_error(bt)
diff --git a/library/python/pytest/yatest_tools.py b/library/python/pytest/yatest_tools.py
index bdcd0cc0b1..0c3138a9fa 100644
--- a/library/python/pytest/yatest_tools.py
+++ b/library/python/pytest/yatest_tools.py
@@ -385,3 +385,37 @@ def _unify_path(path):
         raise MissingTestModule("Can't find proper module for '{}' path among: {}".format(path, suff_tree))
     else:
         return path
+
+
+def colorize_pytest_error(text):
+    error_prefix = "E   "
+    blocks = [text]
+
+    while True:
+        text = blocks.pop()
+
+        err_start = text.find(error_prefix, 1)
+        if err_start == -1:
+            return ''.join(blocks + [text])
+
+        for pos in range(err_start + 1, len(text) - 1):
+            if text[pos] == '\n':
+                if not text[pos + 1:].startswith(error_prefix):
+                    err_end = pos + 1
+                    break
+        else:
+            err_end = len(text)
+
+        bt, error, tail = text[:err_start], text[err_start:err_end], text[err_end:]
+
+        filters = [
+            # File path, line number and function name
+            (re.compile(r"^(.*?):(\d+): in (\S+)", flags=re.MULTILINE),
+             r"[[unimp]]\1[[rst]]:[[alt2]]\2[[rst]]: in [[alt1]]\3[[rst]]"),
+        ]
+        for regex, substitution in filters:
+            bt = regex.sub(substitution, bt)
+
+        blocks.append(bt)
+        blocks.append('[[bad]]' + error)
+        blocks.append(tail)