一日一技:Python類型標注的高級用法
假設你正在寫后端代碼,其中一個函數的功能是傳入文章id,返回文章詳情。因為項目比較大,因此在定義函數時,把類型標注加上,標明了參數的類型和返回的類型。例如:
from typing import List
from dataclasses import dataclass
@dataclass
class ArticleDetail:
id: int
title: str
content: str
tag: List[str]
def query_article_detail(article_id: int) -> ArticleDetail:
detail = ArticleDetail(
id=article_id,
title='文章標題',
content='文章內容',
tag=['tag1', 'tag2']
)
return detail
def test_query_article_detail():
detail = query_article_detail(123)
print(detail.content)
現在,當你拿到返回的detail變量時,IDE的自動補全就可以正常工作了,如下圖所示。
圖片
你想讓這個函數支持批量查詢文章詳情的功能,代碼類似這樣:
def query_article_detail(article_id: int | List[int]) -> ArticleDetail | List[ArticleDetail]:
if isinstance(article_id, int):
detail = ArticleDetail(
id=article_id,
title='文章標題',
cnotallow='文章內容',
tag=['tag1', 'tag2']
)
return detail
else:
details = []
for _id in article_id:
detail = ArticleDetail(
id=_id,
title='文章標題',
cnotallow='文章內容',
tag=['tag1', 'tag2']
)
details.append(detail)
return details
如果傳入的參數是int類型的文章id,那么就返回這篇文章的詳情ArticleDetail對象。如果傳入的是文章列表,那么就返回ArticleDetail對象列表。
現在問題來了,由于query_article_detail函數返回的數據類型不同,如何讓IDE的自動補全能夠正確提示呢?例如當我們傳入了一個文章id列表,但是卻直接讀取返回數據的.content屬性,在IDE上面看不出任何問題,如下圖所示。但顯然會報錯,因為此時的detail變量的值是一個列表。列表是沒有.content屬性的。
圖片
有沒有什么辦法能夠讓IDE根據query_article_detail參數的類型,提示我們對返回數據的使用是否正確呢?
這個場景下,就可以使用Python的typing模塊中的@overload裝飾器,實現函數重載來提示。示例代碼如下:
from typing import List, overload
from dataclasses import dataclass
@dataclass
class ArticleDetail:
id: int
title: str
content: str
tag: List[str]
@overload
def query_article_detail(article_id: List[int]) -> List[ArticleDetail]:
...
@overload
def query_article_detail(article_id: int) -> ArticleDetail:
...
def query_article_detail(article_id: int | List[int]) -> ArticleDetail | List[ArticleDetail]:
if isinstance(article_id, int):
detail = ArticleDetail(
id=article_id,
title='文章標題',
cnotallow='文章內容',
tag=['tag1', 'tag2']
)
return detail
else:
details = []
for _id in article_id:
detail = ArticleDetail(
id=_id,
title='文章標題',
cnotallow='文章內容',
tag=['tag1', 'tag2']
)
details.append(detail)
return details
def test_query_article_detail():
detail = query_article_detail([123, 456, 789])
print(detail.)
在定義函數之前,先使用@overload裝飾器,裝飾兩次函數名。每一次使用不同的參數:
@overload
def query_article_detail(article_id: List[int]) -> List[ArticleDetail]:
...
@overload
def query_article_detail(article_id: int) -> ArticleDetail:
...
這兩個函數都是空函數,函數體用三個點代替。當然你也可以使用pass。而你真正的query_article_detail放到最下面。現在,當我們對detail對象使用自動補全時,IDE就能根據參數的類型來補全對應的值了。
當傳入參數是單個id時,如下圖所示:
圖片
當傳入的參數是id列表時,如下圖所示:
圖片
需要注意的時,所有重載的函數與真正執行的函數,函數名必須全部相同,如下圖所示:
圖片
并且,真正實現功能的函數,必須放在重載函數的下面。
使用這種方式,以后即時別的文件導入并使用你這個函數,你也不用擔心它用錯數據類型了。