如何修改Python包装函数的返回类型提示？

我正在用Python实现Rust风格的Result枚举，写了一个result装饰器，用来自动将函数返回的Ok/Error实例包装成Result对象，代码如下：

import functools

def result(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            res = func(*args, **kwargs)
        except Exception as e:
            res = Error(e)
        
        if isinstance(res, Error) or isinstance(res, Ok):
            return Result(res)
        else:
            raise ValueError("@result function returned non-result (Ok, Error) object")
    return wrapper

但使用该装饰器后，IDE（比如VS Code）显示的函数返回类型提示是Ok | Error，而实际返回的是Result实例。例如下面的函数：

@result
def divide(a: int, b: int):
   if a == 0 or b == 0:
        return Error(DivideByZeroError)
   return Ok(a / b)

悬停时显示的签名是divide(int, int) -> Ok | Error，而非正确的divide(int, int) -> Result。

附Ok、Error、Result的实现代码：

from typing import Any, Callable, final

"""
The Result(Error(Exception)) state
"""
@final
class Error:
    def __init__(self, inner: Exception):
        self.inner = inner

"""
The Result(Ok(inner)) state
"""
@final
class Ok:
    def __init__(self, inner: Any):
        self.inner = inner

"""
The python implementation of Rust's Result enum
Has two states: either Ok(result) or Error(exception)
Note that this class can raise exceptions unlike Option
"""
@final
class Result:
    """
    A Rust-like Result class
    Note that you should not create these manually, but instead use the
    @result decorator and return an Ok or Error directly
    """
    def __init__(self, state: Ok | Error):
        self._state = state

    """
    If the state is Ok, returns the inner, else raises the Error
    """
    def unwrap(self):
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            raise self._state.inner
    
    """
    If the state is Ok, returns the inner, else return the provided 'other'
    """
    def unwrap_or(self, other: Any):
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            return other
    
    """
    If the state is Ok, return the inner, else run the provided callable
    """
    def unwrap_or_else(self, other: Callable):
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            return other()
        
    """
    The internal state of the Result, Ok or Error
    """
    @property
    def state(self) -> Ok | Error:
        return self._state

    def __repr__(self) -> str:
        if isinstance(self._state, Ok):
            return f"Result(Ok({self._state.inner}))"
        else:
            return f"Result(Error({self._state.inner}))"

    """
    Turns the result of a function into a Result:
    If it returns a Result, simply pass this on
    If it returns an Ok or Error, wrap this in a Result and return it
    If it raises an error, return a Result(Error(Exception))
    Else, return a Result(Ok(item))
    """
    @classmethod
    def from_function(cls, func: Callable, *args, **kwargs) -> "Result":
        try:
            res = func(*args, **kwargs)
        except Exception as e:
            res = e

        if isinstance(res, Result):
            return res
        elif isinstance(res, Ok) or isinstance(res, Error):
            return Result(res)
        elif isinstance(res, Exception):
            return Result(Error(res))
        else:
            return Result(Ok(res))

要修正装饰器的返回类型提示，需要给装饰器添加精确的类型注解，利用Python的泛型类型工具保留原函数参数信息，并明确指定包装后的返回类型。

1. 基础修正方案 #

导入typing模块的ParamSpec和TypeVar，给装饰器添加类型签名，将原函数的返回类型（Ok | Error）替换为Result：

import functools
from typing import Callable, ParamSpec, TypeVar

# 定义泛型参数：保留原函数的参数类型
P = ParamSpec("P")
# 限制原函数的返回值只能是Ok或Error
R = TypeVar("R", bound=Ok | Error)

def result(func: Callable[P, R]) -> Callable[P, Result]:
    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result:
        try:
            res = func(*args, **kwargs)
        except Exception as e:
            res = Error(e)
        
        if isinstance(res, Error) or isinstance(res, Ok):
            return Result(res)
        else:
            raise ValueError("@result function returned non-result (Ok, Error) object")
    return wrapper

修改后，被装饰的函数会正确显示返回类型为Result。

2. 进阶泛型优化方案 #

如果需要更精确的类型提示（比如明确Result携带的成功值和错误类型），可以给Ok、Error、Result添加泛型支持：

from typing import Generic, TypeVar, final, Callable, ParamSpec

# 定义泛型类型变量：成功值类型
T = TypeVar("T")
# 定义泛型类型变量：错误类型（限制为Exception子类）
E = TypeVar("E", bound=Exception)

@final
class Ok(Generic[T]):
    def __init__(self, inner: T):
        self.inner: T = inner

@final
class Error(Generic[E]):
    def __init__(self, inner: E):
        self.inner: E = inner

@final
class Result(Generic[T, E]):
    def __init__(self, state: Ok[T] | Error[E]):
        self._state: Ok[T] | Error[E] = state

    def unwrap(self) -> T:
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            raise self._state.inner
    
    def unwrap_or(self, other: T) -> T:
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            return other
    
    def unwrap_or_else(self, other: Callable[[], T]) -> T:
        if isinstance(self._state, Ok):
            return self._state.inner
        else:
            return other()
        
    @property
    def state(self) -> Ok[T] | Error[E]:
        return self._state

    def __repr__(self) -> str:
        if isinstance(self._state, Ok):
            return f"Result(Ok({self._state.inner}))"
        else:
            return f"Result(Error({self._state.inner}))"

对应的装饰器也同步更新为泛型版本：

P = ParamSpec("P")
T = TypeVar("T")
E = TypeVar("E", bound=Exception)

def result(func: Callable[P, Ok[T] | Error[E]]) -> Callable[P, Result[T, E]]:
    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, E]:
        try:
            res = func(*args, **kwargs)
        except Exception as e:
            res = Error(e)
        
        if isinstance(res, Error) or isinstance(res, Ok):
            return Result(res)
        else:
            raise ValueError("@result function returned non-result (Ok, Error) object")
    return wrapper

此时可以给业务函数添加精确的返回类型注解：

class DivideByZeroError(Exception):
    pass

@result
def divide(a: int, b: int) -> Ok[float] | Error[DivideByZeroError]:
   if a == 0 or b == 0:
        return Error(DivideByZeroError())
   return Ok(a / b)

IDE会显示该函数的返回类型为Result[float, DivideByZeroError]，类型提示更加精准。

内容的提问来源于stack exchange，提问作者Windows Documentation Sucks