magic_set.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import types
  2. import inspect
  3. # Modifies class methods (or instances of them) on the fly
  4. # http://blog.ianbicking.org/2007/08/08/opening-python-classes/
  5. # http://svn.colorstudy.com/home/ianb/recipes/magicset.py
  6. def magic_set(obj):
  7. """
  8. Adds a function/method to an object. Uses the name of the first
  9. argument as a hint about whether it is a method (``self``), class
  10. method (``cls`` or ``klass``), or static method (anything else).
  11. Works on both instances and classes.
  12. >>> class color:
  13. ... def __init__(self, r, g, b):
  14. ... self.r, self.g, self.b = r, g, b
  15. >>> c = color(0, 1, 0)
  16. >>> c # doctest: +ELLIPSIS
  17. <__main__.color instance at ...>
  18. >>> @magic_set(color)
  19. ... def __repr__(self):
  20. ... return '<color %s %s %s>' % (self.r, self.g, self.b)
  21. >>> c
  22. <color 0 1 0>
  23. >>> @magic_set(color)
  24. ... def red(cls):
  25. ... return cls(1, 0, 0)
  26. >>> color.red()
  27. <color 1 0 0>
  28. >>> c.red()
  29. <color 1 0 0>
  30. >>> @magic_set(color)
  31. ... def name():
  32. ... return 'color'
  33. >>> color.name()
  34. 'color'
  35. >>> @magic_set(c)
  36. ... def name(self):
  37. ... return 'red'
  38. >>> c.name()
  39. 'red'
  40. >>> @magic_set(c)
  41. ... def name(cls):
  42. ... return cls.__name__
  43. >>> c.name()
  44. 'color'
  45. >>> @magic_set(c)
  46. ... def pr(obj):
  47. ... print obj
  48. >>> c.pr(1)
  49. 1
  50. """
  51. def decorator(func):
  52. is_class = (isinstance(obj, type)
  53. or isinstance(obj, types.ClassType))
  54. args, varargs, varkw, defaults = inspect.getargspec(func)
  55. if not args or args[0] not in ('self', 'cls', 'klass'):
  56. # Static function/method
  57. if is_class:
  58. replacement = staticmethod(func)
  59. else:
  60. replacement = func
  61. elif args[0] == 'self':
  62. if is_class:
  63. replacement = func
  64. else:
  65. def replacement(*args, **kw):
  66. return func(obj, *args, **kw)
  67. try:
  68. replacement.func_name = func.func_name
  69. except:
  70. pass
  71. else:
  72. if is_class:
  73. replacement = classmethod(func)
  74. else:
  75. def replacement(*args, **kw):
  76. return func(obj.__class__, *args, **kw)
  77. try:
  78. replacement.func_name = func.func_name
  79. except:
  80. pass
  81. setattr(obj, func.func_name, replacement)
  82. return replacement
  83. return decorator
  84. if __name__ == '__main__':
  85. import doctest
  86. doctest.testmod()