From f4b48d208ee966dfe057abfc42302ad6b8eb89de Mon Sep 17 00:00:00 2001
From: isra17 <isra017@gmail.com>
Date: Wed, 26 Jul 2023 18:39:18 -0400
Subject: [PATCH] traceback: limit locals stringification

---
 src/saturn_engine/utils/traceback_data.py | 53 ++++++++++++++++++++++-
 tests/utils/test_traceback_data.py        | 17 ++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 tests/utils/test_traceback_data.py

diff --git a/src/saturn_engine/utils/traceback_data.py b/src/saturn_engine/utils/traceback_data.py
index f2c0c64c..4966013d 100644
--- a/src/saturn_engine/utils/traceback_data.py
+++ b/src/saturn_engine/utils/traceback_data.py
@@ -6,6 +6,9 @@
 import linecache
 import sys
 import traceback
+from collections.abc import Collection
+from collections.abc import Mapping
+from contextlib import suppress
 from types import FrameType
 from types import TracebackType
 
@@ -20,6 +23,54 @@
 )
 
 
+def format_local(anyval: object, maxlen: int = 80) -> str:
+    val: str = "<???>"
+    with suppress(Exception):
+        if isinstance(anyval, (str, bytes)):
+            val = repr(anyval[:maxlen])
+        elif isinstance(anyval, (int, float)):
+            val = repr(anyval)
+        elif isinstance(anyval, Mapping):
+            val = ""
+            if not isinstance(anyval, dict):
+                val = f"{type(anyval)}"
+            val += "{"
+            vals = []
+            valslen = len(val)
+            for k, v in anyval.items():
+                ks = format_local(k, maxlen=maxlen - valslen)
+                valslen += len(ks)
+                vs = format_local(v, maxlen=maxlen - valslen)
+                valslen += len(vs) + 4
+                vals.append(f"{ks}: {vs}")
+                if valslen > maxlen:
+                    return val + ", ".join(vals)
+            return val + ", ".join(vals) + "}"
+        elif isinstance(anyval, Collection):
+            val = ""
+            if not isinstance(anyval, list):
+                val = f"{type(anyval)}"
+            val += "["
+            vals = []
+            valslen = len(val)
+            for v in anyval:
+                vs = format_local(v, maxlen=maxlen - valslen)
+                valslen += len(vs) + 2
+                vals.append(vs)
+                if valslen > maxlen:
+                    return val + ", ".join(vals)
+            return val + ", ".join(vals) + "]"
+        else:
+            val = str(type(anyval))
+
+    if len(val) >= maxlen:
+        if val[-1] == "'":
+            val = val[:-1] + "<...>"
+        else:
+            val += "<...>"
+    return val
+
+
 @dataclasses.dataclass
 class TracebackData:
     """Very alike traceback.TracebackException, with a few tweak to ensure
@@ -114,7 +165,7 @@ def extract_stack(
             linecache.lazycache(filename, f.f_globals)
             f_locals = f.f_locals
             _locals: dict[str, str] = (
-                {k: str(v) for k, v in f_locals.items()} if f_locals else {}
+                {k: format_local(v) for k, v in f_locals.items()} if f_locals else {}
             )
 
             firstlineno = co.co_firstlineno
diff --git a/tests/utils/test_traceback_data.py b/tests/utils/test_traceback_data.py
new file mode 100644
index 00000000..509e0186
--- /dev/null
+++ b/tests/utils/test_traceback_data.py
@@ -0,0 +1,17 @@
+from saturn_engine.utils.traceback_data import format_local
+
+
+def test_format_local() -> None:
+    assert format_local("123") == "'123'"
+    assert format_local(b"123") == "b'123'"
+    assert format_local(123) == "123"
+    assert format_local(1.23) == "1.23"
+    assert (
+        format_local({"a": 123, "b": (123, 4), "c": [1, 2, {"e": "f"}]})
+        == "{'a': 123, 'b': <class 'tuple'>[123, 4], 'c': [1, 2, {'e': 'f'}]}"
+    )
+
+    assert (
+        format_local({"a": "a", "b": "b" * 80, "c": "c"})
+        == "{'a': 'a', 'b': '" + "b" * 66 + "<...>"
+    )