Pylance none-checking with a function

Question

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.

score 1 · Accepted Answer · answered Mar 18 '22 at 06:42

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 be None
Store inner_model as a private variable and expose it via a property that does the check and returns a SomeModel:

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)

Your second suggestion here is great - I'll probably take that approach! — Robert Jeffrey, Jun 14 '22 at 08:24

score 0 · Answer 2 · answered Mar 18 '22 at 06:15

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.

Pylance none-checking with a function

2 Answers2

Linked