magic_set.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import inspect
  2. import six
  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. # including python 3 fixes for func_name => __name__ and types.ClassType => type
  7. def magic_set(obj):
  8. """
  9. Adds a function/method to an object. Uses the name of the first
  10. argument as a hint about whether it is a method (``self``), class
  11. method (``cls`` or ``klass``), or static method (anything else).
  12. Works on both instances and classes.
  13. >>> class color:
  14. ... def __init__(self, r, g, b):
  15. ... self.r, self.g, self.b = r, g, b
  16. >>> c = color(0, 1, 0)
  17. >>> c # doctest: +ELLIPSIS
  18. <__main__.color instance at ...>
  19. >>> @magic_set(color)
  20. ... def __repr__(self):
  21. ... return '<color %s %s %s>' % (self.r, self.g, self.b)
  22. >>> c
  23. <color 0 1 0>
  24. >>> @magic_set(color)
  25. ... def red(cls):
  26. ... return cls(1, 0, 0)
  27. >>> color.red()
  28. <color 1 0 0>
  29. >>> c.red()
  30. <color 1 0 0>
  31. >>> @magic_set(color)
  32. ... def name():
  33. ... return 'color'
  34. >>> color.name()
  35. 'color'
  36. >>> @magic_set(c)
  37. ... def name(self):
  38. ... return 'red'
  39. >>> c.name()
  40. 'red'
  41. >>> @magic_set(c)
  42. ... def name(cls):
  43. ... return cls.__name__
  44. >>> c.name()
  45. 'color'
  46. >>> @magic_set(c)
  47. ... def pr(obj):
  48. ... print obj
  49. >>> c.pr(1)
  50. 1
  51. """
  52. def decorator(func):
  53. is_class = isinstance(obj, six.class_types)
  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.__name__ = 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.__name__ = func.__name__
  79. except:
  80. pass
  81. setattr(obj, func.__name__, replacement)
  82. return replacement
  83. return decorator
  84. if __name__ == '__main__':
  85. import doctest
  86. doctest.testmod()