Merge pull request #13 from tusharsadhwani/feat/stderr-and-return-code

Add stderr and return code support

tests/test_files/returncode.py

@@ -0,0 +1,6 @@
+_, _, return_code = ~"false"
+assert return_code == 1
+
+exit_code = 123
+_, _, return_code = ~f"exit {exit_code}"
+assert return_code == exit_code

tests/zxpy_test.py

@@ -18,14 +18,33 @@ def test_shell_output(command: str, output: str) -> None:
     assert zx.run_shell(command).rstrip('\r\n') == output
 
 
-def test_files() -> None:
-    test_files = [
-        './tests/test_files/yeses.py',
-    ]
+@pytest.mark.parametrize(
+    ('command', 'stdout', 'stderr', 'return_code'),
+    (
+        ('echo hello world', 'hello world\n', '', 0),
+        ('echo -n failed && exit 200', 'failed', '', 200),
+        ('cat .', '', 'cat: .: Is a directory\n', 1),
+    )
+)
+def test_stdout_stderr_returncode(
+        command: str,
+        stdout: str,
+        stderr: str,
+        return_code: int,
+) -> None:
+    assert zx.run_shell_alternate(command) == (stdout, stderr, return_code)
 
-    for filepath in test_files:
-        filename = os.path.basename(filepath)
 
-        with open(filepath) as file:
-            module = ast.parse(file.read())
-            zx.run_zxpy(filename, module)
+@pytest.mark.parametrize(
+    ('filepath'),
+    (
+        './tests/test_files/yeses.py',
+        './tests/test_files/returncode.py',
+    ),
+)
+def test_files(filepath) -> None:
+    filename = os.path.basename(filepath)
+
+    with open(filepath) as file:
+        module = ast.parse(file.read())
+        zx.run_zxpy(filename, module)

zx.py

@@ -27,7 +27,7 @@
 import subprocess
 import sys
 import traceback
-from typing import Literal, Optional, Union, overload
+from typing import Literal, Optional, Tuple, Union, overload
 
 
 def cli() -> None:
@@ -88,6 +88,26 @@ def run_shell(command: str, print_it: bool = False) -> Optional[str]:
     return None
 
 
+def run_shell_alternate(command: str) -> Tuple[str, str, int]:
+    """Like run_shell but returns 3 values: stdout, stderr and return code"""
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        shell=True,
+    )
+    process.wait()
+    assert process.stdout is not None
+    assert process.stderr is not None
+    assert process.returncode is not None
+
+    return (
+        process.stdout.read().decode(),
+        process.stderr.read().decode(),
+        process.returncode,
+    )
+
+
 def run_zxpy(filename: str, module: ast.Module) -> None:
     """Runs zxpy on a given file"""
     patch_shell_commands(module)
@@ -108,16 +128,25 @@ def patch_shell_commands(module: Union[ast.Module, ast.Interactive]) -> None:
 class ShellRunner(ast.NodeTransformer):
     """Replaces the ~'...' syntax with run_shell(...)"""
     @staticmethod
-    def modify_expr(expr: ast.expr) -> ast.expr:
+    def modify_expr(
+            expr: ast.expr,
+            return_stderr_and_returncode: bool = False,
+    ) -> ast.expr:
         if (
             isinstance(expr, ast.UnaryOp)
             and isinstance(expr.op, ast.Invert)
             and isinstance(expr.operand, (ast.Str, ast.JoinedStr))
         ):
+            function_name = (
+                'run_shell_alternate'
+                if return_stderr_and_returncode
+                else 'run_shell'
+            )
+
             return ast.Call(
-                func=ast.Name(id='run_shell', ctx=ast.Load()),
+                func=ast.Name(id=function_name, ctx=ast.Load()),
                 args=[expr.operand],
-                keywords=[]
+                keywords=[],
             )
 
         return expr
@@ -128,7 +157,11 @@ def visit_Expr(self, expr: ast.Expr) -> ast.Expr:
         return expr
 
     def visit_Assign(self, assign: ast.Assign) -> ast.Assign:
-        assign.value = self.modify_expr(assign.value)
+        assign.value = self.modify_expr(
+            assign.value,
+            return_stderr_and_returncode=isinstance(assign.targets[0], ast.Tuple),
+        )
+
         super().generic_visit(assign)
         return assign