Python TypeHint 一个根据输入 returns 不同类型的函数

Question

我有一个这样工作的方法，它通常用于 return a Path.


from typing import Literal, Union
from pathlib import Path

def get_filename(return_type: Literal["Path", "str"]) -> Union[Path, str]:
    filename = "./foo/bar.csv"

    if return_type == "str":
        return filename
    elif return_type == "Path":
        return Path(filename)
    else:
        raise ValueError(
            f'Return type must be either "Path" or "str", not {return_type}'
        )


file = get_filename(return_type="Path")
print(file.is_file())

在最后一行，我从 Pylance 收到以下消息：

(method) is_file: () -> bool | Unknown Whether this path is a regular file (also True for symlinks pointing to regular files).
Cannot access member "is_file" for type "str" Member "is_file" is unknownPylancereportGeneralTypeIssues

有没有办法正确键入提示这种情况，以便 Pylance 知道 file 是 Path？或者我应该让它总是 return 路径并有另一个方法调用 get_filename 将输出转换为字符串然后 return 那？

谢谢

编辑 1

我刚刚意识到另一个更常见的场景：

import pandas as pd

# returns dataframe
df = pd.read_csv(file)

# returns a series, Pylance doesn't know this
series = pd.read_csv(file, squeeze=True)

这里 pandas 输入参数可以改变输出类型，Pylance 也不能处理这个。要让 Pylance 知道 series 是 pd.Series，您必须执行以下操作：

# return series which pylance is happy with
df = pd.read_csv(file)
series = df.squeeze()

Answer 1

我想你可以使用 typing.overload:

from typing import Literal, overload
from pathlib import Path

@overload
def get_filename(return_type: Literal["Path"]) -> Path:...

@overload
def get_filename(return_type: Literal["str"]) -> str:...

def get_filename(return_type):
    # your code goes here

Python TypeHint 一个根据输入 returns 不同类型的函数

Python TypeHint a function that returns different types based on input

python

type-hinting

pylance