Python TypeHint:根据输入返回不同类型的函数

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得到以下消息:

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

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

感谢编辑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