Python 裝飾器之 functools.wraps

Decorator

在以前的 Blog 中曾經簡單寫過 Decorator。這次需要講的更細一些。

Decorator 通過返回包裝物件實現間接呼叫,以此插入額外邏輯。是從老大那邊偷來的哪裡摘抄來的，應該算是言簡意賅了。

@dec2
@dec1
def func(arg1, arg2, ...):
    pass

1

2

3

4

@dec2

@dec1

def func(arg1, arg2, ...):

pass

可以還原成

def func(arg1, arg2, ...):
    pass
func = dec2(dec1(func))

1

2

3

def func(arg1, arg2, ...):

pass

func = dec2(dec1(func))

@decomaker(argA, argB, ...)
def func(arg1, arg2, ...):
    pass

1

2

3

@decomaker(argA, argB, ...)

def func(arg1, arg2, ...):

pass

可以還原成

func = decomaker(argA, argB, ...)(func)

1	func = decomaker(argA, argB, ...)(func)

In [1]: def outer(func):
   ...:     def inner():
   ...:         print "before func"
   ...:         ret = func()
   ...:         return ret + 1
   ...:     return inner #返回 inner 函式物件
   ...:

In [2]: @outer  # 直譯器執⾏行 foo = outer(foo)
   ...: def foo():
   ...:     return 1
   ...:

In [3]: foo
Out[3]: <function __main__.inner>

In [4]: foo()
        before func
Out[4]: 2

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

In [1]: def outer(func):

...: def inner():

...: print "before func"

...: ret = func()

...: return ret + 1

...: return inner #返回 inner 函式物件

...:

In [2]: @outer # 直譯器執⾏行 foo = outer(foo)

...: def foo():

...: return 1

...:

In [3]: foo

Out[3]: <function __main__.inner>

In [4]: foo()

before func

Out[4]: 2

這個過程中執行了下面幾步

函式 foo 作為裝飾器 outer 的引數被傳入
函式 inner 對 func 進行呼叫，然後裝飾器 outer 返回 inner
原來的函式名 foo 關聯到 inner，如上面的foo 所示，呼叫 foo 時間上是在呼叫 inner

裝飾器不僅可以用函式返回包裝物件，也可以是個類，不過這種方法太尼瑪囉嗦，這裡就不介紹了，想了解的自己去翻吧。下面我們寫一個有點用處的 Decorator。假想我們有個coordinate類，而且這個類提供了 x, y座標，而我們要對兩個coordinate 物件進行計算。程式碼如下：

class Coordinate(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y
    def __repr__(self):
        return "Coord: " + str(self.__dict__)

def add(a, b):
    return Coordinate(a.x + b.x, a.y + b.y)

def sub(a, b):
    return Coordinate(a.x - b.x, a.y - b.y)

In [8]: one = Coordinate(100, 200)

In [9]: two = Coordinate(300, 200)

In [10]: three = Coordinate(-100, -100)

In [11]: sub(one, three)
Out[11]: Coord: {'y': 300, 'x': 200}

In [12]: add(one, three)
Out[12]: Coord: {'y': 100, 'x': 0}

In [13]: sub(one, two)
Out[13]: Coord: {'y': 0, 'x': -200}

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

class Coordinate(object):

def __init__(self, x, y):

self.x = x

self.y = y

def __repr__(self):

return "Coord: " + str(self.__dict__)

def add(a, b):

return Coordinate(a.x + b.x, a.y + b.y)

def sub(a, b):

return Coordinate(a.x - b.x, a.y - b.y)

In [8]: one = Coordinate(100, 200)

In [9]: two = Coordinate(300, 200)

In [10]: three = Coordinate(-100, -100)

In [11]: sub(one, three)

Out[11]: Coord: {'y': 300, 'x': 200}

In [12]: add(one, three)

Out[12]: Coord: {'y': 100, 'x': 0}

In [13]: sub(one, two)

Out[13]: Coord: {'y': 0, 'x': -200}

上面例子中的sub(one, two)與three都有負數，當我們把座標限制在第一象限時，這兩個就不符合我們的要求，用 Decorator 來做一個檢測再好不過了

In [14]: def wrapper(func):
   ....:     def checker(a, b):
   ....:         if a.x < 0 or a.y < 0:
   ....:             a = Coordinate(a.x if a.x > 0 else 0, a.y if a.y > 0 else 0)
   ....:         if b.x < 0 or b.y < 0:
   ....:             b = Coordinate(b.x if b.x > 0 else 0, b.y if b.y > 0 else 0)
   ....:         ret = func(a, b)
   ....:         if ret.x < 0 or ret.y <0:
   ....:             ret = Coordinate(ret.x if ret.x > 0 else 0, ret.y if ret.y > 0 else 0)
   ....:         return ret
   ....:     return checker
   ....:
In [16]: @wrapper
   ....: def add(a, b):
   ....:     return Coordinate(a.x + b.x, a.y + b.y)
   ....:

In [17]: @wrapper
   ....: def sub(a, b):
   ....:     return Coordinate(a.x - b.x, a.y + b.y)
   ....:

In [18]: add(one, three)
Out[18]: Coord: {'y': 200, 'x': 100}

In [19]: one
Out[19]: Coord: {'y': 200, 'x': 100}

In [20]: sub(one, two)
Out[20]: Coord: {'y': 400, 'x': 0}

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

In [14]: def wrapper(func):

....: def checker(a, b):

....: if a.x < 0 or a.y < 0:

....: a = Coordinate(a.x if a.x > 0 else 0, a.y if a.y > 0 else 0)

....: if b.x < 0 or b.y < 0:

....: b = Coordinate(b.x if b.x > 0 else 0, b.y if b.y > 0 else 0)

....: ret = func(a, b)

....: if ret.x < 0 or ret.y <0:

....: ret = Coordinate(ret.x if ret.x > 0 else 0, ret.y if ret.y > 0 else 0)

....: return ret

....: return checker

....:

In [16]: @wrapper

....: def add(a, b):

....: return Coordinate(a.x + b.x, a.y + b.y)

....:

In [17]: @wrapper

....: def sub(a, b):

....: return Coordinate(a.x - b.x, a.y + b.y)

....:

In [18]: add(one, three)

Out[18]: Coord: {'y': 200, 'x': 100}

In [19]: one

Out[19]: Coord: {'y': 200, 'x': 100}

In [20]: sub(one, two)

Out[20]: Coord: {'y': 400, 'x': 0}

這樣，只計算的函式add與sub前面加一個 Decorator 就可以完成座標的校驗。比在函式內實現要優雅一些。

Decorator 還可以為類增加額外的成員，

In [21]: def hello(cls):
   ....:     cls.hello = staticmethod(lambda: "HELLO")
   ....:     return cls
   ....:

In [22]: <a href="http://www.jobbole.com/members/Hello">@hello</a>
   ....: class World(object):pass
   ....:

In [23]: World.hello
Out[23]: <function __main__.<lambda>>

In [24]: World.hello()
Out[24]: 'HELLO'

1

2

3

4

5

6

7

8

9

10

11

12

13

14

In [21]: def hello(cls):

....: cls.hello = staticmethod(lambda: "HELLO")

....: return cls

....:

In [22]: <a href="http://www.jobbole.com/members/Hello">@hello</a>

....: class World(object):pass

....:

In [23]: World.hello

Out[23]: <function __main__.<lambda>>

In [24]: World.hello()

Out[24]: 'HELLO'

functools.wraps

我們在使用 Decorator 的過程中，難免會損失一些原本的功能資訊。直接拿 stackoverflow 裡面的栗子

def logged(func):
    def with_logging(*args, **kwargs):
        print func.__name__ + " was called"
        return func(*args, **kwargs)
    return with_logging

@logged
def f(x):
   """does some math"""
   return x + x * x

def f(x):
    """does some math"""
    return x + x * x
f = logged(f)

In [24]: f.__name__
Out[24]: with_logging

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

def logged(func):

def with_logging(*args, **kwargs):

print func.__name__ + " was called"

return func(*args, **kwargs)

return with_logging

@logged

def f(x):

"""does some math"""

return x + x * x

def f(x):

"""does some math"""

return x + x * x

f = logged(f)

In [24]: f.__name__

Out[24]: with_logging

而functools.wraps 則可以將原函式物件的指定屬性複製給包裝函式物件, 預設有 __module__、__name__、__doc__,或者通過引數選擇。程式碼如下：

from functools import wraps
def logged(func):
    @wraps(func)
    def with_logging(*args, **kwargs):
        print func.__name__ + " was called"
        return func(*args, **kwargs)
    return with_logging

@logged
def f(x):
   """does some math"""
   return x + x * x

print f.__name__  # prints 'f'
print f.__doc__   # prints 'does some math'

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

from functools import wraps

def logged(func):

@wraps(func)

def with_logging(*args, **kwargs):

print func.__name__ + " was called"

return func(*args, **kwargs)

return with_logging

@logged

def f(x):

"""does some math"""

return x + x * x

print f.__name__ # prints 'f'

print f.__doc__ # prints 'does some math'

Python 裝飾器之 functools.wraps

functools.wraps

相關文章