Merge pull request #744 from hartmans/await_return_cmd

await command with _return_cmd True returns RunningCommand

sh.py

@@ -889,7 +889,10 @@ def __next__(self):
     def __await__(self):
         async def wait_for_completion():
             await self.aio_output_complete.wait()
-            return str(self)
+            if self.call_args["return_cmd"]:
+                return self
+            else:
+                return str(self)
 
         return wait_for_completion().__await__()
 

tests/sh_test.py

@@ -1732,7 +1732,7 @@ def test_async_exc(self):
         py = create_tmp_test("""exit(34)""")
 
         async def producer():
-            await python(py.name, _async=True)
+            await python(py.name, _async=True, _return_cmd=False)
 
         self.assertRaises(sh.ErrorReturnCode_34, asyncio.run, producer())
 
@@ -1786,6 +1786,22 @@ async def producer():
 
         self.assertRaises(sh.ErrorReturnCode_34, asyncio.run, producer())
 
+    def test_async_return_cmd(self):
+        py = create_tmp_test(
+            """
+import sys
+sys.exit(0)
+"""
+        )
+
+        async def main():
+            result = await python(py.name, _async=True, _return_cmd=True)
+            self.assertIsInstance(result, sh.RunningCommand)
+            result_str = await python(py.name, _async=True, _return_cmd=False)
+            self.assertIsInstance(result_str, str)
+
+        asyncio.run(main())
+
     def test_handle_both_out_and_err(self):
         py = create_tmp_test(
             """