automain.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. # Copyright 2014-2015 Nathan West
  2. #
  3. # This file is part of autocommand.
  4. #
  5. # autocommand is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU Lesser General Public License as published by
  7. # the Free Software Foundation, either version 3 of the License, or
  8. # (at your option) any later version.
  9. #
  10. # autocommand is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU Lesser General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public License
  16. # along with autocommand. If not, see <http://www.gnu.org/licenses/>.
  17. import sys
  18. from .errors import AutocommandError
  19. class AutomainRequiresModuleError(AutocommandError, TypeError):
  20. pass
  21. def automain(module, *, args=(), kwargs=None):
  22. '''
  23. This decorator automatically invokes a function if the module is being run
  24. as the "__main__" module. Optionally, provide args or kwargs with which to
  25. call the function. If `module` is "__main__", the function is called, and
  26. the program is `sys.exit`ed with the return value. You can also pass `True`
  27. to cause the function to be called unconditionally. If the function is not
  28. called, it is returned unchanged by the decorator.
  29. Usage:
  30. @automain(__name__) # Pass __name__ to check __name__=="__main__"
  31. def main():
  32. ...
  33. If __name__ is "__main__" here, the main function is called, and then
  34. sys.exit called with the return value.
  35. '''
  36. # Check that @automain(...) was called, rather than @automain
  37. if callable(module):
  38. raise AutomainRequiresModuleError(module)
  39. if module == '__main__' or module is True:
  40. if kwargs is None:
  41. kwargs = {}
  42. # Use a function definition instead of a lambda for a neater traceback
  43. def automain_decorator(main):
  44. sys.exit(main(*args, **kwargs))
  45. return automain_decorator
  46. else:
  47. return lambda main: main