@@ -97,6 +97,25 @@ parameter as follows: This will result in a short option of -q and a long option of --custom-name. +Special arguments +----------------- + +Arguments with the names 'stdin', 'stdout' or 'stderr' will be automatically +passed the relevant Python objects, for example: + + #!/usr/bin/env python + # upper.py + import optfunc + + @optfunc.main + def upper_stdin(stdin, stdout): + stdout.write(stdin.read().upper()) + +Does the following: + + $ echo "Hello, world" | ./upper.py + HELLO, WORLD + Subcommands ----------- README.txt @@ -71,10 +71,11 @@ def resolve_args(func, argv): parser, required_args = func_to_optionparser(func) options, args = parser.parse_args(argv) - # Special case for stdin - if 'stdin' in required_args: - required_args.remove('stdin') - options.optfunc_use_stdin = True + # Special case for stdin/stdout/stderr + for pipe in ('stdin', 'stdout', 'stderr'): + if pipe in required_args: + required_args.remove(pipe) + setattr(options, 'optfunc_use_%s' % pipe, True) # Do we have correct number af required args? if len(required_args) != len(args): @@ -90,7 +91,9 @@ def resolve_args(func, argv): return options.__dict__, parser._errors -def run(func, argv=None, stderr=sys.stderr, stdin=sys.stdin): +def run( + func, argv=None, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr + ): argv = argv or sys.argv[1:] include_func_name_in_errors = False @@ -123,9 +126,10 @@ def run(func, argv=None, stderr=sys.stderr, stdin=sys.stdin): else: raise TypeError('arg is not a Python function or class') - # Special case for stdin - if resolved.pop('optfunc_use_stdin', False): - resolved['stdin'] = stdin + # Special case for stdin/stdout/stderr + for pipe in ('stdin', 'stdout', 'stderr'): + if resolved.pop('optfunc_use_%s' % pipe, False): + resolved[pipe] = locals()[pipe] if not errors: try: @@ -141,8 +145,10 @@ def run(func, argv=None, stderr=sys.stderr, stdin=sys.stdin): def main(*args, **kwargs): prev_frame = inspect.stack()[-1][0] - if inspect.getmodule(prev_frame).__name__ == '__main__': + mod = inspect.getmodule(prev_frame) + if mod is not None and mod.__name__ == '__main__': run(*args, **kwargs) + return args[0] # So it won't break anything if used as a decorator # Decorators def notstrict(fn): optfunc.py @@ -234,5 +234,37 @@ class TestOptFunc(unittest.TestCase): optfunc.run(func, stdin=FakeStdin()) self.assertEqual(consumed, ['hello']) + def test_stdout_special_argument(self): + def upper(stdin, stdout): + stdout.write(stdin.read().upper()) + + class FakeStdin(object): + def read(self): + return "hello" + + class FakeStdout(object): + written = '' + def write(self, w): + self.written = w + + stdout = FakeStdout() + self.assertEqual(stdout.written, '') + optfunc.run(upper, stdin=FakeStdin(), stdout=stdout) + self.assertEqual(stdout.written, 'HELLO') + + def test_stderr_special_argument(self): + def upper(stderr): + stderr.write('an error') + + class FakeStderr(object): + written = '' + def write(self, w): + self.written = w + + stderr = FakeStderr() + self.assertEqual(stderr.written, '') + optfunc.run(upper, stderr=stderr) + self.assertEqual(stderr.written, 'an error') + if __name__ == '__main__': unittest.main() test.py eda71154f5a5d8ddab05d752d5ebf8fa896cb8e5 Simon Willison simon@simonwillison.net http://github.com/simonw/optfunc/commit/e3fa034a545ed94ac5a039cf5b170c7d0ee21b7b e3fa034a545ed94ac5a039cf5b170c7d0ee21b7b 2009-05-29T15:47:52-07:00 2009-05-29T15:47:52-07:00 Added support for stdout and stderr special arguments and documented them 674ccc3d2dee422db2bf28d2af09381af3e70391 Simon Willison simon@simonwillison.net