Update heterograph.py (#6506)

Co-authored-by: Hongzhi (Steve), Chen <chenhongzhi.nkcs@gmail.com>

python/dgl/heterograph.py

@@ -6061,7 +6061,9 @@ class DGLGraph(object):
         old_eframes = self._edge_frames
         self._node_frames = [fr.clone() for fr in self._node_frames]
         self._edge_frames = [fr.clone() for fr in self._edge_frames]
+        try:
             yield
+        finally:
             self._node_frames = old_nframes
             self._edge_frames = old_eframes
 

tests/python/common/function/test_basics.py

@@ -677,6 +677,19 @@ def test_local_scope(idtype):
 
     foo(g)
 
+    # test exception handling
+    def foo(g):
+        try:
+            with g.local_scope():
+                g.ndata["hh"] = F.ones((g.num_nodes(), 1))
+                # throw TypeError
+                1 + "1"
+        except TypeError:
+            pass
+        assert "hh" not in g.ndata
+
+    foo(g)
+
 
 @parametrize_idtype
 def test_isolated_nodes(idtype):