[Feature] Address PR review comments

Remove _SAFE_SHADOW_NAMES -- all field names that shadow TensorDict
attributes now require explicit TypedTensorDict["shadow"]. Add Python
version comments to try/except import fallbacks.

Made-with: Cursor

Author: vmoens

tensordict/typedtensordict.py

@@ -14,6 +14,7 @@
 from tensordict.base import NO_DEFAULT
 
 try:
+    # Python 3.11+ (PEP 681)
     from typing import dataclass_transform
 except ImportError:
 
@@ -25,13 +26,11 @@ def _identity(cls):
 
 
 try:
+    # Python 3.11+ (PEP 655)
     from typing import NotRequired
 except ImportError:
     from typing_extensions import NotRequired  # noqa: F401
 
-# Fields on TensorDict that are safe to shadow without the "shadow" option.
-_SAFE_SHADOW_NAMES = frozenset({"_is_non_tensor", "data"})
-
 # Annotation names that are class-level metadata, not user fields.
 _META_FIELDS = frozenset(
     {
@@ -262,7 +261,7 @@ def __new__(
         if not cls._shadow:
             td_dir = _get_td_dir()
             for attr in expected:
-                if attr in td_dir and attr not in _SAFE_SHADOW_NAMES:
+                if attr in td_dir:
                     raise AttributeError(
                         f"Field '{attr}' shadows a TensorDict attribute. "
                         f"Use TypedTensorDict['shadow'] to allow this."
@@ -273,9 +272,9 @@ def __new__(
         cls.__optional_keys__ = optional
 
         # Generate properties for fields that clash with TensorDict attributes
-        # so they override the parent's version. When shadow=False, only
-        # _SAFE_SHADOW_NAMES fields (like "data") get properties; the rest
-        # were already rejected above.
+        # so they override the parent's version. When shadow=False these
+        # fields were already rejected above, so this only runs when
+        # shadow=True.
         td_dir = _get_td_dir()
         for attr in expected:
             if attr in td_dir:

test/test_typedtensordict.py

@@ -520,16 +520,13 @@ class Empty(TypedTensorDict):
         assert e.batch_size == torch.Size([3])
         assert len(e.keys()) == 0
 
-    def test_safe_shadow_data_without_option(self):
-        """'data' is in _SAFE_SHADOW_NAMES so it works without shadow=True."""
-
-        class WithData(TypedTensorDict):
-            data: Tensor
-            x: Tensor
+    def test_data_requires_shadow(self):
+        """'data' shadows a TensorDict attribute and requires shadow=True."""
+        with pytest.raises(AttributeError, match="shadows a TensorDict attribute"):
 
-        wd = WithData(data=torch.randn(3, 2), x=torch.randn(3, 4), batch_size=[3])
-        assert wd.data.shape == (3, 2)
-        assert wd.x.shape == (3, 4)
+            class WithData(TypedTensorDict):
+                data: Tensor
+                x: Tensor
 
     def test_batch_iter(self):
         """__iter__ on TypedTensorDict iterates over batch dimension."""