Skip to content

Commit

Permalink
Updated code_utils.py & local_commandline_code_executor.py (powershel…
Browse files Browse the repository at this point in the history
…l to pwsh) (microsoft#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]>
  • Loading branch information
3 people authored Feb 19, 2024
1 parent 8af7c0f commit e58f3a1
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 5 deletions.
30 changes: 27 additions & 3 deletions autogen/code_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")


Expand Down
2 changes: 1 addition & 1 deletion autogen/coding/local_commandline_code_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
40 changes: 39 additions & 1 deletion test/test_code_utils.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -11,6 +13,7 @@
execute_code,
extract_code,
improve_code,
get_powershell_command,
improve_function,
infer_lang,
is_docker_running,
Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit e58f3a1

Please sign in to comment.