Python TypeHint a function 根据输入返回不同类型

Question

我有一个像这样工作的方法，它通常用于返回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 收到以下消息：

(方法) is_file: () -> bool | 未知此路径是否为常规文件（对于指向常规文件的符号链接也为 True）。
无法访问类型“str”的成员“is_file”成员“is_file”未知PylancereportGeneralTypeIssues

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

谢谢

编辑 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 中，输入参数可以更改 output 类型，而 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