PythonTypeHint基于输入返回不同类型的函数

Question

我有一个这样的方法，它通常用于返回路径。

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：（）-

有没有一种方法可以正确地键入提示这种情况，以便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)

在熊猫中，输入参数可以改变输出类型，Pylance也不能处理这个问题。Pylance知道系列是一个pd。系列你必须做：

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

Answer 1

我想你可以用打字来完成。过载：

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