Update heterograph.py (#6506)

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

python/dgl/heterograph.py

@@ -6061,9 +6061,11 @@ def local_scope(self):
         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]
-        yield
-        self._node_frames = old_nframes
-        self._edge_frames = old_eframes
+        try:
+            yield
+        finally:
+            self._node_frames = old_nframes
+            self._edge_frames = old_eframes
 
     def formats(self, formats=None):
         r"""Get a cloned graph with the specified allowed sparse format(s) or

tests/python/common/function/test_basics.py

@@ -677,6 +677,19 @@ def foo(g):
 
     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):