Python 小技巧 —— 用類寫裝飾器

知識 10-01

程序員大咖

點擊右側關注，免費進階高級！

Fossen,紅塵練心,fossen.cn

地址：zhihu.com/people/all-ming-yun

最近學到了一個有趣的裝飾器寫法，就記錄一下。

裝飾器是一個返回函數的函數。寫一個裝飾器，除了最常見的在函數中定義函數以外，Python還允許使用類來定義一個裝飾器。

1、用類寫裝飾器

下面用常見的寫法實現了一個緩存裝飾器。

def

cache


 
(func):
    
data
 = {}
    def
 wrapper(*args, **kwargs):
        key
 = f"{func.__name__}-{str(args)}-{str(kwargs)})"

        if
 key
 in
 data
:
            result
 = data.get(key
 
)
            print(
"cached"
)
        else
:
            result
 = func(*args, **kwargs)
            data
[key
] = result

            print("calculated"
)
        return
 result

    return

wrapper

看看緩存的效果。

@cache


def
 rectangle_area
(length, width)
:

    return
 length * width
rectangle_area(
2
, 3
)
# calculated
 
# 6

rectangle_area(
2
, 3
)
# cached

# 6

裝飾器的@cache是一個語法糖，相當於func = cache(func)，如果這裡的cache不是一個函數，而是一個類又會怎樣呢？定義一個類class Cache，那麼調用func = Cache(func)會得到一個對象，這時返回的func其實是Cache的對象。定義__call__方法可以將類的實例變成可調用對象，可以像調用函數一樣調用對象。然後在__call__方法里調用原本的func函數就能實現裝飾器了。所以Cache類也能當作裝飾器使用，並且能以@Cache的形式使用。

接下來把cache函數改寫為Cache類：

class

Cache

:
    def
 __init__(self
 
, func):
        self.func = func
        self.data = {}
    
def
 __call__(self
, *args, **kwargs):
        func = self.func
        data
 = self.data
        key
 = f"{func.__name__}-{str(args)}-{str(kwargs)})"

        if
 key
 in
 data
:
            result
 
 = data.get(
key
)
            print("cached"
)
        else
:
            result
 = func(*args, **kwargs)
            data
[key
] = result

            print("calculated"
)
        return
 result

再看看緩存結果，效果一樣。

@Cache


def
 rectangle_area
(length, width)
:

    return
 length * width
rectangle_area(
2
, 3
)
# calculated

# 6

rectangle_area(
2
, 3
)
# cached

# 6

2、裝飾類的方法

裝飾器不止能裝飾函數，也經常用來裝飾類的方法，但是我發現用類寫的裝飾器不能直接用在裝飾類的方法上。（有點繞…）

先看看函數寫的裝飾器如何裝飾類的方法。

class

 Rectangle
:

    
def
 __init__
(
self
, length, width)
:
        self
.length = length
        self
.width = width
    @cache
    
def
 area
(
self
)
:
        return
 self
.length * self
.width
r = Rectangle(
2
, 3
)
r.area()
# calculated

# 6

r.area()
# cached

# 6

但是如果直接換成Cache類會報錯，這個錯誤的原因是area被裝飾後變成了類的一個屬性，而不是方法。

class

 Rectangle
:

    
def
 __init__
(
self
, length, width)
:
        self
.length = length
        self
.width = width
    @Cache
    
def
 area
(
self
)
:
        return
 self
.length * self
.width
r = Rectangle(
2
, 3
)
r.area()
# TypeError: area() missing 1 required positional argument: "self"

Rectangle.area
# <__main__.Cache object at 0x0000012D8E7A6D30>

r.area
# <__main__.Cache object at 0x0000012D8E7A6D30>

回頭再來看看沒有裝飾器的情況，Python在實例化對象後把函數變成了方法。

class

 Rectangle
:

    
def
 __init__
(
self
, length, width)
:
        self
.length = length
        self
.width = width
    
def
 area
(
self
)
:
        return
 self
.length * self
.width
Rectangle.area
# <function Rectangle.area at 0x0000012D8E7B28C8>

r = Rectangle(2
, 3
)
r.area
# <bound method Rectangle.area of <__main__.Rectangle object

因此解決辦法很簡單，要用類寫的裝飾器來裝飾類的方法，只需要把可調用對象包裝成函數就行。

# 定義一個簡單的裝飾器，什麼也不做，僅僅是把可調用對象包裝成函數


def
 method
(call)
:
    
def
 wrapper
(*args, **kwargs)
:
        return
 call(*args, **kwargs)
    return
 wrapper
class
 Rectangle
:

    
def
 __init__
(
self
, length, width)
:
        self
.length = length
        self
.width = width
    @method
    @Cache
    
def
 area
(
self
)
:
        return
 self
.length * self
.width
r = Rectangle(
2
, 3
)
r.area()
# calculated

# 6

r.area()
# cached

# 6

或者用@property還能直接把方法變成屬性。

class

 Rectangle
:

    
def
 __init__
(
self
, length, width)
:
        self
.length = length
        self
.width = width
    @property
    @Cache
    
def
 area
(
self
)
:
        return
 self
.length * self
.width
r = Rectangle(
2
, 3
)
r.area
# calculated

# 6

r.area
# cached

# 6