Wednesday, December 31, 2014

Defining __str__ for class in Python

From time to time I use classes as enums:
class POLL_BASE(object): pass
class POLL_READ(POLL_BASE): pass
class POLL_WRITE(POLL_BASE): pass
No I can inspect whether function returns a final value or a hint that I need to poll its socket to wait for more data.
rv = work()
if issubclass(rv, p.POLL_BASE):
    print "work() returned %s" % rv
    # poll socket
else:
    # we are done!
Now while debugging it you'll notice that print statement above will output this:
work() returned <class '__main__.POLL_WRITE'>
The above is readable, but I wanted something more pretty. However trying to define __str__ method on POLL_BASE class did not help - in fact it will be invoked on POLL_BASE instances but not POLL_BASE itself.

Here is the trick: since __str__ is invoked on instances, then we need to define it on instance class. The "class" of POLL_BASE is type so we can define our own metaclass for POLL_BASE:

class POLL_META(type):
    # This is for fancy printing of POLL_* class names
    def __str__(cls):
        return cls.__name__

class POLL_BASE(object):
    __metaclass__ = POLL_META
Now our fancy printing will work:
work() returned POLL_WRITE