diff --git a/pyuac/__init__.py b/pyuac/__init__.py index 8f214fc..0eaf1c6 100644 --- a/pyuac/__init__.py +++ b/pyuac/__init__.py @@ -21,7 +21,7 @@ See also this utility function which runs a function as admin and captures the stdout/stderr: -run_function_as_admin_with_output(my_main_function) +run_function_as_admin(my_main_function) """ diff --git a/pyuac/run_function.py b/pyuac/run_function.py index c0b7aa9..5e1c82b 100644 --- a/pyuac/run_function.py +++ b/pyuac/run_function.py @@ -3,7 +3,7 @@ # vim: fileencoding=utf-8 tabstop=4 expandtab shiftwidth=4 """ -See run_function_as_admin_with_output +See run_function_as_admin TODO: this is unfinished and untested """ @@ -17,7 +17,7 @@ from pyuac import isUserAdmin, runAsAdmin -def run_function_as_admin_with_output( +def run_function_as_admin( run_function, run_args=None, run_kwargs=None, return_output=False, stdout_handle=None, stderr_handle=None, scan_for_error=True, @@ -44,12 +44,13 @@ def run_function_as_admin_with_output( @param stdout_handle: file handle to write the process stdout output, defaults to sys.stdout @param stderr_handle: file handle to write the process stderr output, defaults to sys.stderr @param scan_for_error: look at the last line only for the string 'error' and - turn that into a RuntimeError if found. + turn that into a RuntimeError if found. ONLY scans the LAST line of stderr and stdout! @param stdout_temp_fn: the name of the temporary log file to use (will be deleted) for standard output stream of the sub-process. If not given, a default is generated @param stderr_temp_fn: the name of the temporary log file to use (will be deleted) for standard error stream of the sub-process. If not given, a default is generated - @return: None unless return_output is set + @return: None unless return_output is set. + If return_output is True, the output is a 2-tuple of (stdout, stderr) strings. """ # Should we add another function parameter to run the in the "not-admin" case? @@ -80,21 +81,32 @@ def run_function_as_admin_with_output( traceback.print_exc(file=stderr_handle) else: runAsAdmin(wait=True) - # TODO: add stderr handling - if os.path.exists(stdout_temp_fn): - with open(stdout_temp_fn, "r") as log_fh: + + rv = [] + for filename, handle in ( + (stdout_temp_fn, stdout_handle), + (stderr_temp_fn, stderr_handle), + ): + if os.path.exists(filename): + with open(filename, "r") as log_fh: console_output = log_fh.read() - os.remove(stdout_temp_fn) - if os.path.exists(stdout_temp_fn): - print("Warning, can't delete " + stdout_temp_fn) + os.remove(filename) + if os.path.exists(filename): + print("Warning, can't delete " + filename) if scan_for_error: - last_line = str.splitlines(console_output.strip())[-1].strip() - if last_line.lower().find("error") != -1: - raise RuntimeError(last_line) + lines = str.splitlines(console_output.strip()) + if lines: + last_line = lines[-1].strip() + if last_line.lower().find("error") != -1: + raise RuntimeError(last_line) if return_output: - return console_output + # return console_output + rv.append(console_output) + + handle.write(console_output) + handle.flush() - stdout_handle.write(console_output) - stdout_handle.flush() + if return_output and rv: + return rv diff --git a/tests/test_run_function.py b/tests/test_run_function.py new file mode 100644 index 0000000..11f8a42 --- /dev/null +++ b/tests/test_run_function.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python +# -*- coding: utf-8; mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vim: fileencoding=utf-8 tabstop=4 expandtab shiftwidth=4 + +""" + +""" + +from pyuac import isUserAdmin +from pyuac.run_function import run_function_as_admin + + +def sample_function_body(arg1, kwarg2='Default'): + print("Hello, world.") + is_admin = isUserAdmin() + print("isUserAdmin: %s" % (is_admin,)) + print("arg1: %s" % (arg1,)) + print("kwarg2: %s" % (kwarg2,)) + return + + +def test_run_function(): + expected_output = """Hello, world. +isUserAdmin: True +arg1: foobar +kwarg2: biz +""" + actual_stdout, actual_stderr = run_function_as_admin( + sample_function_body, + ('foobar',), + {'kwarg2': 'biz'}, + return_output=True + ) + assert actual_stdout == expected_output + assert actual_stderr == "" + return