fix(bt_json): serialize nested values in logged dataclasses

py/src/braintrust/bt_json.py

@@ -3,6 +3,7 @@
 import math
 import warnings
 from collections.abc import Callable, Mapping
+from functools import cache
 from typing import Any, NamedTuple, cast, overload
 
 
@@ -16,9 +17,30 @@
     _HAS_ORJSON = False
 
 
-def _to_bt_safe(v: Any) -> Any:
+class _LoggerTypes(NamedTuple):
+    Span: type
+    Experiment: type
+    Dataset: type
+    Logger: type
+    BaseAttachment: type
+    ReadonlyAttachment: type
+
+
+@cache
+def _logger_types() -> _LoggerTypes:
+    # Imported lazily to avoid circular imports, and cached because `_to_bt_safe` runs per logged value.
+    from braintrust.logger import BaseAttachment, Dataset, Experiment, Logger, ReadonlyAttachment, Span
+
+    return _LoggerTypes(Span, Experiment, Dataset, Logger, BaseAttachment, ReadonlyAttachment)
+
+
+def _to_bt_safe(v: Any, copy_dataclass_field: Callable[[Any, int], Any] | None = None, depth: int = 0) -> Any:
     """
     Converts the object to a Braintrust-safe representation (i.e. Attachment objects are safe (specially handled by background logger)).
+
+    `copy_dataclass_field(value, depth)` converts dataclass field values. `bt_safe_deep_copy` passes its own
+    walker and the fields' depth so they share its cycle and depth tracking; without it, the dataclass
+    starts a new `bt_safe_deep_copy`.
     """
     # Fast path: check primitives via type identity before hitting
     # isinstance checks against abstract classes or Pydantic model_dump.
@@ -37,8 +59,7 @@ def _to_bt_safe(v: Any) -> Any:
     if isinstance(v, (int, str)):
         return v
 
-    # avoid circular imports
-    from braintrust.logger import BaseAttachment, Dataset, Experiment, Logger, ReadonlyAttachment, Span
+    Span, Experiment, Dataset, Logger, BaseAttachment, ReadonlyAttachment = _logger_types()
 
     if isinstance(v, Span):
         return "<span>"
@@ -59,10 +80,9 @@ def _to_bt_safe(v: Any) -> Any:
         return v.reference
 
     if dataclasses.is_dataclass(v) and not isinstance(v, type):
-        # Use manual field iteration instead of dataclasses.asdict() because
-        # asdict() deep-copies values, which breaks objects like Attachment
-        # that contain non-copyable items (thread locks, file handles, etc.)
-        return {f.name: _to_bt_safe(getattr(v, f.name)) for f in dataclasses.fields(v)}
+        if copy_dataclass_field is None:
+            return bt_safe_deep_copy(v)
+        return _dataclass_to_dict(v, copy_dataclass_field, depth)
 
     # Pydantic model classes (not instances) with model_json_schema
     if isinstance(v, type) and hasattr(v, "model_json_schema") and callable(cast(Any, v).model_json_schema):
@@ -100,6 +120,20 @@ def _to_bt_safe(v: Any) -> Any:
     return bt_loads(bt_dumps(v, encoder=_str_encoder))
 
 
+def _dataclass_to_dict(v: Any, copy_field: Callable[[Any, int], Any], depth: int) -> dict[str, Any]:
+    # Use manual field iteration instead of dataclasses.asdict() because
+    # asdict() deep-copies values, which breaks objects like Attachment
+    # that contain non-copyable items (thread locks, file handles, etc.)
+    result = {}
+    for f in dataclasses.fields(v):
+        try:
+            result[f.name] = copy_field(getattr(v, f.name), depth)
+        except Exception:
+            # Keep the other fields (and any attachments in them) when one can't be read.
+            result[f.name] = f"<non-sanitizable: {type(v).__name__}.{f.name}>"
+    return result
+
+
 @overload
 def bt_safe_deep_copy(
     obj: Mapping[str, Any],
@@ -155,43 +189,14 @@ def _deep_copy_object(v: Any, depth: int = 0) -> Any:
         if depth >= max_depth:
             return "<max depth exceeded>"
 
-        # Fast path for dict (the most common container in log data).
-        # Uses type identity instead of isinstance(v, Mapping) which is slow.
-        if tv is dict:
-            obj_id = id(v)
-            if obj_id in visited:
-                return "<circular reference>"
-            visited_add(obj_id)
-            try:
-                result = {}
-                for k in v:
-                    if type(k) is str:
-                        key_str = k
-                    else:
-                        try:
-                            key_str = str(k)
-                        except Exception:
-                            key_str = f"<non-stringifiable-key: {type(k).__name__}>"
-                    result[key_str] = _deep_copy_object(v[k], depth + 1)
-                return result
-            finally:
-                visited_discard(obj_id)
-        elif tv is list or tv is tuple:
-            obj_id = id(v)
-            if obj_id in visited:
-                return "<circular reference>"
-            visited_add(obj_id)
-            try:
-                return [_deep_copy_object(x, depth + 1) for x in v]
-            finally:
-                visited_discard(obj_id)
-        # Slow path for non-builtin Mapping/set types.
-        elif isinstance(v, Mapping):
-            obj_id = id(v)
-            if obj_id in visited:
-                return "<circular reference>"
-            visited_add(obj_id)
-            try:
+        obj_id = id(v)
+        if obj_id in visited:
+            return "<circular reference>"
+        visited_add(obj_id)
+        try:
+            # Fast paths for builtin containers (the most common values in log data).
+            # Type identity checks avoid the slow isinstance(v, Mapping) for them.
+            if tv is dict or (tv is not list and tv is not tuple and isinstance(v, Mapping)):
                 result = {}
                 for k in v:
                     if type(k) is str:
@@ -203,22 +208,17 @@ def _deep_copy_object(v: Any, depth: int = 0) -> Any:
                             key_str = f"<non-stringifiable-key: {type(k).__name__}>"
                     result[key_str] = _deep_copy_object(v[k], depth + 1)
                 return result
-            finally:
-                visited_discard(obj_id)
-        elif isinstance(v, set):
-            obj_id = id(v)
-            if obj_id in visited:
-                return "<circular reference>"
-            visited_add(obj_id)
-            try:
+            if tv is list or tv is tuple or isinstance(v, set):
                 return [_deep_copy_object(x, depth + 1) for x in v]
-            finally:
-                visited_discard(obj_id)
 
-        try:
-            return _to_bt_safe(v)
-        except Exception:
-            return f"<non-sanitizable: {type(v).__name__}>"
+            try:
+                # Dataclass fields go back through this walker so nested attachments keep
+                # their identity and share its cycle/depth tracking.
+                return _to_bt_safe(v, _deep_copy_object, depth + 1)
+            except Exception:
+                return f"<non-sanitizable: {type(v).__name__}>"
+        finally:
+            visited_discard(obj_id)
 
     return _deep_copy_object(obj)
 
@@ -234,15 +234,12 @@ def _to_json_safe(obj: Any) -> Any:
     """
     Handler for non-JSON-serializable objects. Returns a string representation of the object.
     """
-    # avoid circular imports
-    from braintrust.logger import BaseAttachment
-
     try:
         v = _to_bt_safe(obj)
 
         # JSON-safe representation of Attachment objects are their reference.
         # If we get this object at this point, we have to assume someone has already uploaded the attachment!
-        if isinstance(v, BaseAttachment):
+        if isinstance(v, _logger_types().BaseAttachment):
             v = v.reference
 
         return v

py/src/braintrust/test_bt_json.py

@@ -1,14 +1,120 @@
 # pyright: reportUnknownVariableType=false
 # pyright: reportUnknownArgumentType=false
 # pyright: reportPrivateUsage=false
+import dataclasses
 import json
 import warnings
 from typing import Any
 from unittest import TestCase
 
+import pytest
 from braintrust.bt_json import bt_dumps, bt_safe_deep_copy
-from braintrust.logger import Attachment, ExternalAttachment
-from braintrust.test_helpers import with_memory_logger  # noqa: F401
+from braintrust.logger import Attachment, ExternalAttachment, ReadonlyAttachment
+from braintrust.test_helpers import DataclassAttachment, with_memory_logger  # noqa: F401
+
+
+@dataclasses.dataclass
+class _Payload:
+    value: Any
+
+
+@dataclasses.dataclass
+class _DataclassReadonlyAttachment(ReadonlyAttachment):
+    reference: dict[str, Any]
+
+    def __post_init__(self):
+        super().__init__(self.reference)
+
+
+def test_dataclass_nested_attachments():
+    attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain")
+    payload = _Payload(
+        {
+            "direct": attachment,
+            "list": [attachment],
+            "tuple": (attachment,),
+            "dataclass": _Payload(attachment),
+            "dataclasses": [_Payload(attachment)],
+        }
+    )
+
+    assert bt_safe_deep_copy(payload) == {
+        "value": {
+            "direct": attachment,
+            "list": [attachment],
+            "tuple": [attachment],
+            "dataclass": {"value": attachment},
+            "dataclasses": [{"value": attachment}],
+        }
+    }
+
+
+@pytest.mark.parametrize(
+    "base_name,placeholder",
+    [("_NoopSpan", "<span>"), ("Experiment", "<experiment>"), ("Dataset", "<dataset>"), ("Logger", "<logger>")],
+)
+def test_dataclass_sdk_object_subclasses_keep_placeholder(base_name, placeholder):
+    from braintrust import logger
+
+    # The `secret` field must never be copied into logged events.
+    obj = dataclasses.make_dataclass(
+        "DataclassObject", [("secret", str, "token")], bases=(getattr(logger, base_name),)
+    )()
+
+    assert bt_safe_deep_copy({"direct": obj, "nested": _Payload([obj])}) == {
+        "direct": placeholder,
+        "nested": {"value": [placeholder]},
+    }
+
+
+def test_dataclass_attachment_subclasses_keep_special_handling():
+    attachment = DataclassAttachment("test-key")
+    reference = {
+        "type": "braintrust_attachment",
+        "key": "readonly-key",
+        "filename": "a.txt",
+        "content_type": "text/plain",
+    }
+    readonly = _DataclassReadonlyAttachment(reference)
+
+    assert bt_safe_deep_copy({"direct": [attachment, readonly], "nested": _Payload([attachment, readonly])}) == {
+        "direct": [attachment, reference],
+        "nested": {"value": [attachment, reference]},
+    }
+
+
+def test_dataclass_cycles_and_depth_share_walker_state():
+    attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain")
+    payload = _Payload([attachment])
+    payload.value.append(payload)
+
+    # Shared (non-circular) references are copied each time; only true cycles are replaced.
+    assert bt_safe_deep_copy([payload, payload]) == [
+        {"value": [attachment, "<circular reference>"]},
+        {"value": [attachment, "<circular reference>"]},
+    ]
+    assert bt_safe_deep_copy(_Payload([_Payload("leaf")]), max_depth=2) == {"value": ["<max depth exceeded>"]}
+
+
+def test_dataclass_unreadable_field_keeps_readable_fields():
+    @dataclasses.dataclass
+    class Payload:
+        files: list[Any]
+        value: int = 0
+        unset: int = dataclasses.field(init=False)
+
+        def __getattribute__(self, name):
+            if name == "value":
+                raise RuntimeError("Cannot read field")
+            return super().__getattribute__(name)
+
+    attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain")
+
+    assert bt_safe_deep_copy(Payload([attachment])) == {
+        "files": [attachment],
+        "value": "<non-sanitizable: Payload.value>",
+        "unset": "<non-sanitizable: Payload.unset>",
+    }
 
 
 class TestBTJson(TestCase):
@@ -586,36 +692,6 @@ class ModelWithAttachment(BaseModel):
         except ImportError:
             self.skipTest("Pydantic not available")
 
-    def test_bt_safe_deep_copy_dataclass_with_attachments(self):
-        """Test that dataclasses with attachments are handled correctly."""
-        from dataclasses import dataclass
-
-        attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain")
-
-        @dataclass
-        class DataWithAttachment:
-            name: str
-            value: int
-            file: Attachment
-
-        data = DataWithAttachment(name="test", value=42, file=attachment)
-
-        result = bt_safe_deep_copy({"data": data})
-
-        # Dataclasses with Attachment fields are now properly converted by
-        # recursively applying _to_bt_safe to each field instead of using
-        # dataclasses.asdict() which would try to deepcopy the Attachment.
-        self.assertIsInstance(result["data"], dict)
-        self.assertEqual(result["data"]["name"], "test")
-        self.assertEqual(result["data"]["value"], 42)
-        self.assertIs(result["data"]["file"], attachment)
-
-        # Attachments in regular dicts also work fine
-        dict_with_attachment = {"name": "test", "value": 42, "file": attachment}
-        result2 = bt_safe_deep_copy({"data": dict_with_attachment})
-        self.assertIsInstance(result2["data"], dict)
-        self.assertIs(result2["data"]["file"], attachment)
-
     def test_bt_safe_deep_copy_circular_in_pydantic_deferred_to_bt_dumps(self):
         """Test that circular references inside Pydantic model results bypass bt_safe_deep_copy detection.
 

py/src/braintrust/test_helpers.py

@@ -1,9 +1,10 @@
+import dataclasses
 import os
 from contextlib import contextmanager
 
 import pytest
 from braintrust import logger
-from braintrust.logger import ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata
+from braintrust.logger import BaseAttachment, ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata
 from braintrust.util import LazyValue
 
 
@@ -129,6 +130,32 @@ def find_span_by_name(spans, name):
     raise AssertionError(f"Expected span named {name!r}. Available spans: {available_names}")
 
 
+@dataclasses.dataclass
+class DataclassAttachment(BaseAttachment):
+    """A `BaseAttachment` subclass that is also a dataclass, for serialization tests."""
+
+    key: str
+
+    @property
+    def reference(self):
+        return {
+            "type": "braintrust_attachment",
+            "key": self.key,
+            "filename": "file.txt",
+            "content_type": "text/plain",
+        }
+
+    @property
+    def data(self):
+        return b"data"
+
+    def upload(self):
+        raise AssertionError("Tests must not upload attachments")
+
+    def debug_info(self):
+        return {"key": self.key}
+
+
 def init_test_exp(experiment_name: str, project_name: str = None):
     """
     Initialize an experiment for testing with fake project and experiment metadata.