Friday, January 21, 2011

custom string class

Here is how to make a subclass of string which has all the same operations available, but which returns instances of that subclass:

import inspect
import functools

class MyStr(str):
    def __radd__(self, other): 
        return MyStr(str.__add__(other, self))
    def __repr__(self): 
        return self.__class__.__name__ + str.__repr__(self)

def _mystr_wrap(f):
    @functools.wraps(f, ('__name__','__doc__'))
    def g(*a, **kw):
        res = f(*a, **kw)
        if res.__class__ == str:
            return MyStr(res)
        return res
    return g

for name, f in inspect.getmembers(str, callable):
    #skip some special cases
    if name not in ['__class__', '__new__', '__str__', '__init__', '__repr__']:
        setattr(MyStr, name, _mystr_wrap(f))

>>> MyStr("cat")
>>> MyStr("cat")+"dog"
>>> MyStr("cat")[1:]

Note: if you find yourself actually doing this, it is probably a bad code smell :-)

No comments:

Post a Comment