Assertions prune the type-inference tree.

This means that you can write

    x = somethingReturningObject

    assert isinstance(x, T)

and the compiler will infer that 'x' has 'T' if you pass the assertion.

Author: braxtonmckee

typed_python/compiler/function_conversion_context.py

@@ -2116,6 +2116,10 @@ def _convert_statement_ast(self, ast, variableStates: FunctionStackState, contro
                         if msgExpr is not None:
                             expr_context.pushException(AssertionError, msgExpr)
 
+            if not definitelyFails:
+                # we are implicitly the 'True' case of this condition
+                self.restrictByCondition(variableStates, ast.test, result=True)
+
             return expr_context.finalize(None, exceptionsTakeFrom=ast), not definitelyFails
 
         if ast.matches.FunctionDef:

typed_python/compiler/tests/type_inference_test.py

@@ -410,3 +410,24 @@ def f(x: OneOf(None, str)):
                 return localVariableTypesKnownToCompiler()
 
         assert f("hi")['x'] == str
+
+    def test_assertion_prunes_types(self):
+        @Entrypoint
+        def f(x: OneOf(None, str)):
+            assert x is not None
+
+            return x
+
+        assert f.resultTypeFor(str).typeRepresentation is str
+        assert f.resultTypeFor(type(None)).typeRepresentation is str
+
+    def test_branch_with_raise_prunes_types(self):
+        @Entrypoint
+        def f(x: OneOf(None, str)):
+            if x is None:
+                raise Exception("FAIL")
+
+            return x
+
+        assert f.resultTypeFor(str).typeRepresentation is str
+        assert f.resultTypeFor(type(None)).typeRepresentation is str