I am currently working on an implementation of the Rust Result enum in Python. To be able to just return the classes Ok and Error, I have a wrapper that will turn these into a Result, as shown here:
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
However, when wrapping a function with this, it shows its return type as Ok | Error even though this is incorrect and it does return a Result. I found ways to change the function parameters’ signature online, however I want to change the hinted return type. How could I do this?
When hovering over this function in VS Code:
@result
def divide(a: int, b: int):
if a == 0 or b == 0:
return Error(DivideByZeroError)
return Ok(a / b)
I get the signature divide(int, int) -> Ok | Error instead of the correct divide(int, int) -> Result.
For clarification, here are the implementations of Ok, Error and Result:
"""
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\n
Has two states: either Ok(result) or Error(exception)\n
Note that this class can raise exceptions unlike Option
"""
@final
class Result:
"""
A Rust-like Result class\n
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:\n
If it returns a Result, simply pass this on\n
If it returns an Ok or Error, wrap this in a Result and return it\n
If it raises an error, return a Result(Error(Exception))\n
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))
>Solution :
Because you didn’t provide any annotations in the result decorator, Pylance assumes that it doesn’t change the signature of the function.
You’ll need to specify the annotations for result. I don’t know all the specifics of your classes, but it should look something like
from collections.abc import Callable
from typing import ParamSpec
P = ParamSpec("P")
def result(func: Callable[P, Ok | Error]) -> Callable[P, Result]:
...