[BC] Add an interactive only mode for env.py

compiler_opt/rl/env.py

@@ -293,17 +293,19 @@ def _get_clang_generator(
     interactive_only: bool = False,
 ) -> Generator[Optional[Tuple[ClangProcess, InteractiveClang]],
                Optional[corpus.LoadedModuleSpec], None]:
-  """Returns a generator for creating InteractiveClang objects.
-
-  TODO: fix this docstring
+  """Returns a tuple of generators for creating InteractiveClang objects.
 
   Args:
     clang_path: Path to the clang binary to use within InteractiveClang.
     task_type: Type of the MLGO task to use.
-    interactive_only: If set to true only iclang is yielded
+    interactive_only: If set to true the returned tuple of generators is
+      iclang, iclang instead of iclang, clang
 
   Returns:
-    The generator for InteractiveClang objects.
+    A tuple of generators created with clang_session. First argument of
+    the tuple is always an interactive clang session. The second argumnet
+    is a default clang session if interactive_only is False and otherwise
+    the exact same interactive clang session object as the first argument.
   """
   while True:
     # The following line should be type-hinted as follows:

compiler_opt/rl/env_test.py

@@ -214,6 +214,7 @@ def test_env(self, mock_popen):
 
       step = test_env.step(np.array([1], dtype=np.int64))
       self.assertEqual(step.step_type, env.StepType.LAST)
+      self.assertNotEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
 
   @mock.patch('subprocess.Popen')
   def test_env_interactive_only(self, mock_popen):
@@ -239,7 +240,7 @@ def test_env_interactive_only(self, mock_popen):
 
       step = test_env.step(np.array([1], dtype=np.int64))
       self.assertEqual(step.step_type, env.StepType.LAST)
-      self.assertEqual(step.reward, {'default': 0.})
+      self.assertEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
 
 
 if __name__ == '__main__':