[mypyc] Use correct rtype for variable with inferred optional type

mypyc/irbuild/builder.py

@@ -48,6 +48,7 @@
 )
 from mypy.types import (
     AnyType,
+    DeletedType,
     Instance,
     ProperType,
     TupleType,
@@ -570,21 +571,22 @@ def get_assignment_target(
                 self.error("Cannot assign to the first argument of classmethod", line)
             if lvalue.kind == LDEF:
                 if symbol not in self.symtables[-1]:
+                    if isinstance(symbol, Var) and not isinstance(symbol.type, DeletedType):
+                        reg_type = self.type_to_rtype(symbol.type)
+                    else:
+                        reg_type = self.node_type(lvalue)
                     # If the function is a generator function, then first define a new variable
                     # in the current function's environment class. Next, define a target that
                     # refers to the newly defined variable in that environment class. Add the
                     # target to the table containing class environment variables, as well as the
                     # current environment.
                     if self.fn_info.is_generator:
                         return self.add_var_to_env_class(
-                            symbol,
-                            self.node_type(lvalue),
-                            self.fn_info.generator_class,
-                            reassign=False,
+                            symbol, reg_type, self.fn_info.generator_class, reassign=False
                         )
 
                     # Otherwise define a new local variable.
-                    return self.add_local_reg(symbol, self.node_type(lvalue))
+                    return self.add_local_reg(symbol, reg_type)
                 else:
                     # Assign to a previously defined variable.
                     return self.lookup(symbol)

mypyc/test-data/run-misc.test

@@ -84,6 +84,36 @@ assert f(a) is a
 assert g(None) == 1
 assert g(a) == 2
 
+[case testInferredOptionalAssignment]
+from typing import Any, Generator
+
+def f(b: bool) -> Any:
+    if b:
+        x = None
+    else:
+        x = 1
+
+    if b:
+        y = 1
+    else:
+        y = None
+
+    m = 1 if b else None
+    n = None if b else 1
+    return ((x, y), (m, n))
+
+def gen(b: bool) -> Generator[Any, None, None]:
+    if b:
+        y = 1
+    else:
+        y = None
+    yield y
+
+assert f(False) == ((1, None), (None, 1))
+assert f(True) == ((None, 1), (1, None))
+assert next(gen(False)) is None
+assert next(gen(True)) == 1
+
 [case testWith]
 from typing import Any
 class Thing: