Updated code_utils.py & local_commandline_code_executor.py (powershell to pwsh) (#1710)

* Update code_utils.py

Updated the powershell command to pwsh

* Update code_utils.py

added a split to handle powershell in the first condition as well

* Update local_commandline_code_executor.py

added "pwsh" as a command option in lang variable

* Update autogen/coding/local_commandline_code_executor.py

Co-authored-by: Eric Zhu <[email protected]>

* Update code_utils.py

* Update code_utils.py

fixed formatting

* Update code_utils.py

defined a function to detect whether 'powershell' or 'pwsh' works and accordingly use the one that works

* Update code_utils.py

fixed formatting

* Update and rename test_code.py to test_code_utils.py

added a unit test for get_powershell_command function in code_utils.py

* Update test_code_utils.py

fixed formatting

* Update test_code_utils.py

fixed formatting

* Update autogen/code_utils.py

Co-authored-by: Chi Wang <[email protected]>

---------

Co-authored-by: Eric Zhu <[email protected]>
Co-authored-by: Chi Wang <[email protected]>

Author: 3 people

autogen/code_utils.py

@@ -213,13 +213,37 @@ def timeout_handler(signum, frame):
     raise TimeoutError("Timed out!")
 
 
+def get_powershell_command():
+    try:
+        result = subprocess.run(["powershell", "$PSVersionTable.PSVersion.Major"], capture_output=True, text=True)
+        if result.returncode == 0:
+            return "powershell"
+
+    except FileNotFoundError:
+        # This means that 'powershell' command is not found so now we try looking for 'pwsh'
+        try:
+            result = subprocess.run(
+                ["pwsh", "-Command", "$PSVersionTable.PSVersion.Major"], capture_output=True, text=True
+            )
+            if result.returncode == 0:
+                return "pwsh"
+
+        except FileNotFoundError:
+            print("Neither powershell nor pwsh is installed.")
+            return None
+
+
+powershell_command = get_powershell_command()
+
+
 def _cmd(lang):
-    if lang.startswith("python") or lang in ["bash", "sh", "powershell"]:
+    if lang.startswith("python") or lang in ["bash", "sh", powershell_command]:
         return lang
     if lang in ["shell"]:
         return "sh"
-    if lang in ["ps1"]:
-        return "powershell"
+    if lang in ["ps1", "pwsh", "powershell"]:
+        return powershell_command
+
     raise NotImplementedError(f"{lang} not recognized in code execution")
 
 

autogen/coding/local_commandline_code_executor.py

@@ -156,7 +156,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeRe
             )
             filename_uuid = uuid.uuid4().hex
             filename = None
-            if lang in ["bash", "shell", "sh"]:
+            if lang in ["bash", "shell", "sh", "pwsh", "powershell", "ps1"]:
                 filename = f"{filename_uuid}.{lang}"
                 exitcode, logs, _ = execute_code(
                     code=code,

test/test_code_utils.py

@@ -1,8 +1,10 @@
 import os
 import tempfile
 import unittest
-
+from unittest.mock import patch
+import sys
 import pytest
+from io import StringIO
 
 import autogen
 from autogen.code_utils import (
@@ -11,6 +13,7 @@
     execute_code,
     extract_code,
     improve_code,
+    get_powershell_command,
     improve_function,
     infer_lang,
     is_docker_running,
@@ -553,6 +556,41 @@ def test_non_dict_in_list(self):
             content_str(content)
 
 
+class TestGetPowerShellCommand(unittest.TestCase):
+    @patch("subprocess.run")
+    def test_get_powershell_command_powershell(self, mock_subprocess_run):
+        # Set up the mock to return a successful result for 'powershell'
+        mock_subprocess_run.return_value.returncode = 0
+        mock_subprocess_run.return_value.stdout = StringIO("5")
+
+        self.assertEqual(get_powershell_command(), "powershell")
+
+    @patch("subprocess.run")
+    def test_get_powershell_command_pwsh(self, mock_subprocess_run):
+        # Set up the mock to return a successful result for 'pwsh'
+        mock_subprocess_run.side_effect = [FileNotFoundError, mock_subprocess_run.return_value]
+        mock_subprocess_run.return_value.returncode = 0
+        mock_subprocess_run.return_value.stdout = StringIO("7")
+
+        self.assertEqual(get_powershell_command(), "pwsh")
+
+    @patch("subprocess.run")
+    def test_get_powershell_command_no_shell(self, mock_subprocess_run):
+        # Set up the mock to simulate 'powershell' and 'pwsh' not found
+        mock_subprocess_run.side_effect = [FileNotFoundError, FileNotFoundError]
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            get_powershell_command()
+            self.assertEqual(fake_out.getvalue().strip(), "Neither powershell nor pwsh is installed.")
+
+    @patch("subprocess.run")
+    def test_get_powershell_command_no_shell_no_output(self, mock_subprocess_run):
+        # Set up the mock to simulate 'powershell' and 'pwsh' not found without printing error message
+        mock_subprocess_run.side_effect = [FileNotFoundError, FileNotFoundError]
+
+        self.assertIsNone(get_powershell_command())
+
+
 if __name__ == "__main__":
     # test_infer_lang()
     test_extract_code()