test and tweak get_strings

oughtinc · Nov 30, 2022 · 212f00d · 212f00d
1 parent a859fcb
commit 212f00d
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 5 deletions.
diff --git a/ice/trace.py b/ice/trace.py
@@ -280,18 +280,20 @@ def get_strings(value) -> list[str]:
     Represent the given value as a short list of short strings
     that can be stored directly in the central trace file and loaded eagerly in the UI.
     """
+    if isinstance(value, dict) and "value" in value:
+        value = value["value"]
+
     if isinstance(value, dict):
-        if "value" in value:
-            value = value["value"]
-        else:
-            value = {k: v for k, v in value.items() if k not in ("self", "record")}
+        value = {k: v for k, v in value.items() if k not in ("self", "record")}
 
     result = get_first_descendant(value)
 
     if isinstance(result, tuple):
         result = list(result)
     if not (result and isinstance(result, list)):
-        result = [f"{result or '()'}"]
+        if result in (None, (), "", [], {}):
+            result = '()'
+        result = [str(result)]
 
     result = get_short_list(result)
     result = [get_short_string(v) for v in result]

diff --git a/tests/test_trace.py b/tests/test_trace.py
@@ -17,3 +17,66 @@ async def test_trace_blocks():
     assert trace.emit_block("quux") == [1, 1]
 
     assert (current_trace.dir / "block_0.jsonl").read_text() == f'"foo"\n"bar"\n"{long}"\nend'
+
+
+def test_get_strings():
+    assert trace.get_strings("foo") == ["foo"]
+    assert trace.get_strings(["foo", "bar"]) == ["foo", "bar"]
+    assert trace.get_strings({"foo": "bar"}) == ["bar"]
+
+    # First descendant list of strings
+    assert trace.get_strings({"foo": {"spam": ["bar", "baz"], "x": "y"}}) == [
+        "bar",
+        "baz",
+    ]
+
+    # Empty values that get converted to '()'
+    assert trace.get_strings([]) == ["()"]
+    assert trace.get_strings(()) == ["()"]
+    assert trace.get_strings({}) == ["()"]
+    assert trace.get_strings(None) == ["()"]
+    assert trace.get_strings("") == ["()"]
+
+    # Other falsy values are just their string representation
+    assert trace.get_strings(0) == ["0"]
+    assert trace.get_strings(0.0) == ["0.0"]
+    assert trace.get_strings(False) == ["False"]
+
+    # First descendant isn't a string, gets converted
+    assert trace.get_strings([1, 2, 3]) == ["1"]
+
+    # Long lists get truncated
+    assert trace.get_strings(["1", "2", "3", "4", "5"]) == ["1", "2", "3", "..."]
+
+    # Long strings get truncated
+    assert trace.get_strings("a" * 100) == ["a" * 35 + "..."]
+
+    # `value` gets extracted
+    assert trace.get_strings({"a": "b", "c": "d"}) == ["b"]
+    assert trace.get_strings({"value": "b", "c": "d"}) == ["b"]
+    assert trace.get_strings({"a": "b", "value": "d"}) == ["d"]
+    assert trace.get_strings({"a": "b", "value": {}}) == ["()"]
+
+    # self and record are omitted
+    assert trace.get_strings({"self": "foo", "record": "bar", "a": "b", "c": "d"}) == [
+        "b"
+    ]
+    assert trace.get_strings({"self": "foo", "record": "bar"}) == ["()"]
+    assert (
+        trace.get_strings(
+            {
+                "x": "y",
+                "value": {"self": "foo", "record": "bar", "a": "b", "c": "d"},
+            }
+        )
+        == ["b"]
+    )
+
+    # Non-strings are filtered out
+    assert trace.get_strings(["a", "b", 3, None, "c"]) == ["a", "b", "c"]
+
+    # Tuples are converted to lists
+    assert trace.get_strings(("a", "b")) == ["a", "b"]
+
+    # Nested lists are not flattened
+    assert trace.get_strings([["a", "b"], ["c", "d"]]) == ["a", "b"]