# Copyright 2014-2015 Nathan West # # This file is part of autocommand. # # autocommand is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # autocommand is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU Lesser General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with autocommand. If not, see . import sys from .errors import AutocommandError class AutomainRequiresModuleError(AutocommandError, TypeError): pass def automain(module, *, args=(), kwargs=None): ''' This decorator automatically invokes a function if the module is being run as the "__main__" module. Optionally, provide args or kwargs with which to call the function. If `module` is "__main__", the function is called, and the program is `sys.exit`ed with the return value. You can also pass `True` to cause the function to be called unconditionally. If the function is not called, it is returned unchanged by the decorator. Usage: @automain(__name__) # Pass __name__ to check __name__=="__main__" def main(): ... If __name__ is "__main__" here, the main function is called, and then sys.exit called with the return value. ''' # Check that @automain(...) was called, rather than @automain if callable(module): raise AutomainRequiresModuleError(module) if module == '__main__' or module is True: if kwargs is None: kwargs = {} # Use a function definition instead of a lambda for a neater traceback def automain_decorator(main): sys.exit(main(*args, **kwargs)) return automain_decorator else: return lambda main: main