一日一技：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='文章標題',
            content='文章內容',
            tag=['tag1', 'tag2']
        )
        return detail
    else:
        details = []
        for _id in article_id:
            detail = ArticleDetail(
                id=_id,
                title='文章標題',
                content='文章內容',
                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='文章標題',
              content='文章內容',
              tag=['tag1', 'tag2']
          )
          return detail
      else:
          details = []
          for _id in article_id:
              detail = ArticleDetail(
                  id=_id,
                  title='文章標題',
                  content='文章內容',
                  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 列表時，如下圖所示：

需要注意的時，所有重載的函數與真正執行的函數，函數名必須全部相同，如下圖所示：

並且，真正實現功能的函數，必須放在重載函數的下面。

使用這種方式，以後即時別的文件導入並使用你這個函數，你也不用擔心它用錯數據類型了。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/wqZJt0tsqh0kqMlIr1uPcQ