Fix Inspector losing debug handles when given an ETRecord object

devtools/inspector/_inspector.py

@@ -1132,15 +1132,28 @@ def _consume_etrecord(self, reference_graph_name) -> None:
             return
 
         # (1) Debug Handle Symbolification
+        # These maps are keyed by instruction id: ints on an ETRecord object
+        # from get_etrecord(), strs once it went through save() and
+        # parse_etrecord() (JSON). _gen_resolve_debug_handles looks up strs.
+        def with_str_keys(mapping: Optional[Dict]) -> Optional[Dict]:
+            if mapping is None:
+                return None
+            return {str(key): value for key, value in mapping.items()}
+
+        debug_handle_map = with_str_keys(self._etrecord._debug_handle_map[FORWARD])
+        delegate_map = (
+            with_str_keys(self._etrecord._delegate_map[FORWARD])
+            if self._etrecord._delegate_map is not None
+            else None
+        )
+        instruction_id_to_num_outs_map = with_str_keys(
+            self._etrecord._instruction_id_to_num_outs_map[FORWARD]
+        )
         for event_block in self.event_blocks:
             event_block._gen_resolve_debug_handles(
-                self._etrecord._debug_handle_map[FORWARD],
-                (
-                    self._etrecord._delegate_map[FORWARD]
-                    if self._etrecord._delegate_map is not None
-                    else None
-                ),
-                self._etrecord._instruction_id_to_num_outs_map[FORWARD],
+                debug_handle_map,
+                delegate_map,
+                instruction_id_to_num_outs_map,
             )
 
         # (2) Event Metadata Association

devtools/inspector/tests/inspector_test.py

@@ -390,6 +390,45 @@ def test_inspector_get_exported_program(self):
                     )
                 )
 
+    def test_inspector_with_etrecord_object_resolves_debug_handles(self):
+        _, _, et_output = TestETRecord().get_test_model(generate_etrecord=True)
+        etrecord = et_output.get_etrecord()
+        # An ETRecord object from get_etrecord() keys its maps by int instruction id
+        instruction_id = next(iter(etrecord._debug_handle_map["forward"]))
+        self.assertIsInstance(instruction_id, int)
+
+        def gen_event_blocks(*args, **kwargs):
+            return [
+                EventBlock(
+                    name=EVENT_BLOCK_NAME,
+                    events=[Event(name="op", _instruction_id=instruction_id)],
+                )
+            ]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            etrecord_path = os.path.join(tmp_dir, "etrecord.bin")
+            etrecord.save(etrecord_path)
+
+            with patch.object(
+                _inspector, "gen_etdump_object", return_value=None
+            ), patch.object(
+                EventBlock, "_gen_from_etdump", side_effect=gen_event_blocks
+            ):
+                inspector_from_object = Inspector(
+                    etdump_path=ETDUMP_PATH, etrecord=etrecord
+                )
+                inspector_from_file = Inspector(
+                    etdump_path=ETDUMP_PATH, etrecord=etrecord_path
+                )
+
+        event_from_object = inspector_from_object.event_blocks[0].events[0]
+        event_from_file = inspector_from_file.event_blocks[0].events[0]
+        self.assertIsNotNone(event_from_file.debug_handles)
+        # Passing the ETRecord object directly must give the same result as
+        # passing the path of the saved ETRecord
+        self.assertEqual(event_from_object.debug_handles, event_from_file.debug_handles)
+        self.assertEqual(event_from_object.op_types, event_from_file.op_types)
+
     def test_populate_debugging_related_fields_raises_for_inconsistent_events(self):
         ret_event: Event = Event(
             name="event",