hacks.py

   1 class Category(type):
   2         """
   3         Adds attributes to an existing class.
   4
   5         """
   6
   7         replace_attrs = False
   8         dunder_attrs = False
   9         never_attrs = ('__module__', '__metaclass__')
  10
  11         def __new__(cls, name, bases, attrs):
  12                 if len(bases) != 1:
  13                         raise AttributeError('%s: "%s" cannot add methods to more than one class.' % (cls.__name__, name))
  14
  15                 base = bases[0]
  16
  17                 for attr, value in attrs.iteritems():
  18                         if attr in cls.never_attrs:
  19                                 continue
  20                         if not cls.dunder_attrs and attr.startswith('__'):
  21                                 continue
  22                         if not cls.replace_attrs and hasattr(base, attr):
  23                                 continue
  24                         setattr(base, attr, value)
  25
  26                 return base
  27
  28
  29 class MonkeyPatch(type):
  30         """
  31         Similar to Category, except it will replace attributes.
  32
  33         """
  34
  35         replace_attrs = True
  36         dunder_attrs = Category.dunder_attrs
  37         never_attrs = Category.never_attrs
  38
  39         unpatches = {}
  40
  41         @classmethod
  42         def unpatched(cls, klass, name):
  43                 try:
  44                         return cls.unpatches[klass][name]
  45                 except:
  46                         return getattr(klass, name)
  47
  48         def __new__(cls, name, bases, attrs):
  49                 if len(bases) != 1:
  50                         raise AttributeError('%s: "%s" cannot patch more than one class.' % (cls.__name__, name))
  51
  52                 base = bases[0]
  53
  54                 for attr, value in attrs.iteritems():
  55                         if attr in cls.never_attrs:
  56                                 continue
  57                         if not cls.dunder_attrs and attr.startswith('__'):
  58                                 continue
  59                         if hasattr(base, attr):
  60                                 if not cls.replace_attrs:
  61                                         continue
  62                                 else:
  63                                         if base not in cls.unpatches:
  64                                                 cls.unpatches[base] = {}
  65                                         cls.unpatches[base][attr] = getattr(base, attr)
  66
  67                         setattr(base, attr, value)
  68
  69                 return base