I have a class that requires a sequence of actions to be taken:
class SomeModel:
def __init__(self):
pass
def predict(self, X):
return None
class Model:
def __init__(self):
self.inner_model = None
def _check_trained(self):
# Simplified; real version has more checks.
assert self.inner_model is not None
def train(self, X, y):
self.inner_model = SomeModel()
# More code here
def predict(self, X):
self._check_trained()
return self.inner_model.predict(X)
Pylance gives a type error on self.inner_model.predict
as self.inner_model
could be None
. However, this is prevented by the previous check. Unrolling the train check function would fix this but would be unwieldly.
Is there a way to verify the not-None property using a function? Or will I need to explicitly disable the type check for this line?
Additional context: VSCode 1.65.2 Pylance v2022.3.2 Pylance on basic Type Checking Mode.
Edit: Here's a screenshot of the behaviour I'm getting.
CodePudding user response:
It doesn't look like Pylance treats the assertion as a postcondition of the function. I couldn't find a specific source for this, but I have seen behaviors like this from other static analysis tools -- automatically generating such postconditions could make the analysis quickly intractable.
You can either use # type: ignore
on the line where the call to .predict()
is made, but if you're accessing the same variable many times over, it would be better to use cast()
, like so:
from typing import cast
def predict(self, X):
self._check_trained()
inner_model = cast(SomeModel, self.inner_model)
return inner_model.predict(X)
Some further discussion on handling nullable values with Pylance can be found here.
CodePudding user response:
Like @brokenbenchmark says, type checking doesn't extend across function boundaries, so treats self.inner_model
as an Optional[Model]
.
One way around this is to return the inner model from self._check_trained()
and typehint its return type.
class Model:
def __init__(self):
self.inner_model = None
def _check_trained(self) -> SomeModel:
# Simplified; real version has more checks.
assert self.inner_model is not None
return self.inner_model
def train(self, X, y):
self.inner_model = SomeModel()
# More code here
def predict(self, X):
inner_model = self._check_trained()
return inner_model.predict(X)
Other possible alternatives:
- enforce assignment of the inner model in
__init__
so it can't beNone
- Store
inner_model
as a private variable and expose it via a property that does the check and returns aSomeModel
:
class Model:
def __init__(self):
self._inner_model = None
@property
def inner_model(self) -> SomeModel:
assert self._inner_model is not None
return self._inner_model
def _check_trained(self):
# Simplified; real version has more checks.
assert self.inner_model is not None
def train(self, X, y):
self._inner_model = SomeModel()
# More code here
def predict(self, X):
self._check_trained()
return self.inner_model.predict(X)