Python 3.8中类型标注触发TypeError的非升级修复方案

问题重现 #

在Python 3.8的Colab环境中运行以下代码：

from typing import List, Dict

def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> list[float]:
    result = openai.Embedding.create(
      model=model,
      input=text
    )
    return result["data"][0]["embedding"]

def compute_doc_embeddings(df: pd.DataFrame) -> dict[tuple[str, str], list[float]]:
    """
    Create an embedding for each row in the dataframe using the OpenAI Embeddings API.
    
    Return a dictionary that maps between each embedding vector and the index of the row that it corresponds to.
    """
    return {
        idx: get_embedding(r.content) for idx, r in df.iterrows()
    }

出现错误：

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-32-bd6d3a14202c> in <module>
      1 from typing import List, Dict
      2 
----> 3 def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> list[float]:
      4     result = openai.Embedding.create(
      5       model=model,

TypeError: 'type' object is not subscriptable

修复方案（无需升级Python 3.9） #

方案1：使用typing模块的泛型类型替代原生类型下标写法 #

Python 3.8及更早版本不支持list[float]、dict[K, V]这类PEP 585引入的原生泛型标注语法，需要改用typing模块提供的List、Dict、Tuple类型：

修改后的代码：

from typing import List, Dict, Tuple
import pandas as pd
import openai

# 需提前定义EMBEDDING_MODEL，示例值如下
EMBEDDING_MODEL = "text-embedding-ada-002"

def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> List[float]:
    result = openai.Embedding.create(
      model=model,
      input=text
    )
    return result["data"][0]["embedding"]

def compute_doc_embeddings(df: pd.DataFrame) -> Dict[Tuple[str, str], List[float]]:
    """
    Create an embedding for each row in the dataframe using the OpenAI Embeddings API.
    
    Return a dictionary that maps between each embedding vector and the index of the row that it corresponds to.
    """
    return {
        idx: get_embedding(r.content) for idx, r in df.iterrows()
    }

方案2：启用__future__ annotations延迟类型解析 #

在代码开头导入from __future__ import annotations，Python会将类型标注当作字符串处理，不会在运行时解析泛型下标，从而避免报错：

修改后的代码：

from __future__ import annotations
from typing import List, Dict
import pandas as pd
import openai

EMBEDDING_MODEL = "text-embedding-ada-002"

def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> list[float]:
    result = openai.Embedding.create(
      model=model,
      input=text
    )
    return result["data"][0]["embedding"]

def compute_doc_embeddings(df: pd.DataFrame) -> dict[tuple[str, str], list[float]]:
    """
    Create an embedding for each row in the dataframe using the OpenAI Embeddings API.
    
    Return a dictionary that maps between each embedding vector and the index of the row that it corresponds to.
    """
    return {
        idx: get_embedding(r.content) for idx, r in df.iterrows()
    }

该方案可保留PEP 585的语法风格，同时兼容Python 3.7+版本。

内容的提问来源于stack exchange，提问作者Mark Wagner