Replacing 'if __name__ == __main__' with decorator (was: Double underscores -- ugly?)

[grflanagan at gmail.com]

On Feb 19, 10:01 am, Duncan Booth <duncan.bo... at invalid.invalid>
wrote:
> Berwyn <berh... at gmail.com> wrote:
> >> Is it just me that thinks "__init__" is rather ugly? Not to mention
> >> "if __name__ == '__main__': ..."?
>
> > That ugliness has long been my biggest bugbear with python, too.  The
> > __name__ == '__main__' thing is something I always have to look up,
> > every time I use it, too ... awkward.
>
> > I'd settle for:
>
> >     hidden def init(self):          # which could be extended to work
> > for everything "hidden x=3"
> >     ...
>
> > And for __name__ == '__main__' how about:
>
> >     if sys.main():
> >         ...
>
> Or even:
>
>   @hidden
>   def init(self): ...
>
> @main
> def mymainfunc():
>    ...
>
> The first of those probably wants some metaclass support to make it work
> cleanly, but here's a sample implementation for the second one:
>
> import sys, atexit
> def main(f):
>     """Decorator for main function"""
>     def runner():
>         sys.exit(f())
>     if f.func_globals['__name__']=='__main__':
>         atexit.register(runner)
>     return f
>
> print "define mymainfunc"
> @main
> def mymainfunc(args=sys.argv):
>     print "Got args", args
>     return 3
> print "end of script"
>
> If you have multiple functions marked as main that will run them in
> reverse order, so it might be better to put them on a list and use a
> single runner to clear the list. Also, I have no idea what happens to
> the exit code if you use this decorator more than once.
>
> BTW, should anyone be wondering, you can still use atexit inside a
> function called from atexit and any registered functions are then called
> when the first one returns.

I liked the idea, thanks. But I also wanted to be able to pass options
to the decorator, for example optparser.OptionParser instances, so
(after wrestling with nested functions for half a day!) I ended up
with the class below, FWIW.

Getting the interplay between __new__, __init__ and __call__ was
somewhat trial-and-error, and maybe I just painted myself into a
corner, but (slightly academic question) is it the case that if you
want to be able to use a decorator both with and without options, ie.
like this:

@mainmethod
def main(...)

and like this:

@mainmethod(parser=myparser)
def main(...)

then you cannot use that decorator for a function that expects or
allows a function as its first argument? Because how and when can you
decide whether that function is the decorated one or the function
parameter?

Anyway, thanks again.

[code]

class mainmethod(object):

    _parser = None
    kwargs = {}
    args =()

    def _getparser(self):
        if self._parser is None:
            self._parser = CommonOptionParser()
        return self._parser
    parser = property(_getparser)

    def __new__(cls, *args, **kw):
        obj = super(mainmethod, cls).__new__(cls, *args, **kw)
        if len(args) == 1 and inspect.isfunction(args[0]):
            #we assume that the decorator has been declared with no
arguments,
            #so go to straight to __call__, don't need __init__
            #if it's the case that the wrapped 'main' method allows or
            #expects a function as its first (and only) positional
argument
            #then you can't use this decorator
            return obj(args[0])
        else:
            return obj

    def __init__(self, *args, **kw):
        self.args = args
        self._parser = kw.pop('parser', None)
        self.kwargs = kw

    def _updatekwargs(self, dict):
        #don't want default null values of parser to overwrite
anything
        #passed to `mainmethod` itself
        for k,v in dict.iteritems():
            #can't do 'if v: ...' because empty configobj evaluates
False
            if v is None or v == '':
                continue
            self.kwargs[k] = v

    def exit(self):
        try:
            log.end()
        except:
            pass

    def run(self):
        options, args = self.parser.parse_args()
        #the following so that command line options are made available
        #to the decorated function as **kwargs
        self._updatekwargs(self.parser.values.__dict__)
        logargs = (
                self.kwargs.get(OPTLOGFILE, None),
                self.kwargs.get(OPTLOGDIR, None),
                self.kwargs.get(OPTLOGPREFIX, ''),
                )
        self.kwargs[OPTLOGFILE] = logstart(*logargs)
        log.info("SCRIPT: " + sys.argv[0])
        conf = self.kwargs.get(OPTCONFIG, None)
        if conf:
            log.info("%s: %s" % (OPTCONFIG.upper(), conf.filename))
        for k,v in self.kwargs.iteritems():
            if v and k not in COMMONOPTS:
                log.info("%s = %s" % (k, v))
        log.divider()
        return sys.exit(self.func(*self.args, **self.kwargs))

    def __call__(self, f):
        if f.func_globals['__name__'] == '__main__':
            self.func = f
            import atexit
            atexit.register(self.exit)
            atexit.register(self.run)
        return f

[/code]

Gerard