Декорируем разом все методы класса в питоне

07.01.2014
@LEXXX_NF

Язык программирования Python славится своими декораторами, которые можно очень легко применить к любой функции или методу.

Для тех, кто забыл или не знает, что такое декораторы, есть небольшая статья в Википедии: http://ru.wikipedia.org/wiki/Декоратор_(шаблон_проектирования).

А для тех, кто хочет узнать, чем отличаются питоновские декораторы от того, что написано в википедии, есть немного запутанная статья на Хабре: http://habrahabr.ru/post/141411/.

Но что если мы хотим применить один и тот же декоратор сразу ко всем методам класса? Например, у нас есть класс-контроллер, методы которого печатают строки, и нам нужно их перехватить и сделать в них некоторые замены. Можно конечно унаследовать класс и перегрузить все методы родительского класса, но есть способ проще!

Представим, что у нас есть такой класс с двумя публичными методами и одним приватным:

class GoodClass:

	def printA(self):
		print 'This is A'

	def printB(self):
		print 'This is B'

	def _printC(self):
		print 'This is C'

Как видно все методы сразу выводят результат с помощью функции print, а нам бы хотелось сначала записать этот вывод в переменную, чтобы потом иметь возможность его обработать.

Сделаем декоратор, который будет перехватывать вывод исходной функции.

import sys, cStringIO
def return_deco(fn):
	def return_wrapped(*args, **kwargs):
		oldstdout = sys.stdout
		try:
			sys.stdout = cStringIO.StringIO()
			fn(*args, **kwargs)
			res = sys.stdout.getvalue()[:-1]
		finally:
			sys.stdout.close()
			sys.stdout = oldstdout
		return res
	return return_wrapped

Теперь можно вручную применить этот декоратор ко всем методам класса, но мы пойдём другим путём и сделаем это автоматически с помощью еще одного декоратора, который будет применяться не к методу, а уже целиком к классу. Отмечу, что декорировать мы будем только публичные методы.

def for_all_methods(decorator):
	def decorate(cls):
		for attr in cls.__bases__[0].__dict__:
			if callable(getattr(cls, attr)) and not attr.startswith('_'):
				setattr(cls, attr, decorator(getattr(cls, attr)))
		return cls
	return decorate

А теперь сделаем класс-потомок, у которого будут все те же методы, что и у родителя, но с применённым декоратором.

@for_all_methods(return_deco)
class BetterClass(GoodClass):
	pass

Посмотрим что получилось. Для этого создадим 2 объекта: исходного класса и декорированного класса.

GoodObject = GoodClass()
BetterObject = BetterClass()

Методы исходного класса сразу печатают строки.

GoodObject.printA()  #  This is A
GoodObject.printB()  #  This is B
GoodObject._printC()  #  This is C

А методы дочернего класса возвращают результат, так что мы его можем потом обработать.

a = BetterObject.printA()
print a.replace(' is', ' was')  #  This was A
b = BetterObject.printB()
print b.replace(' is', ' was')  #  This was B

При этом приватный метод остался без изменений.

BetterObject._printC()  #  This is C

Так можно не только перехватывать вывод всех методов, но и замерить время их выполнения или сменить формат вывода со строк на JSON и т. д. и т. п.

Комментов нет совсем... почему-то...

Писáть здесь