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

…15206)

```python
def f(b: bool) -> None:
    if b:
        y = 1
    else:
        y = None

f(False)
```

On encountering y = 1, mypyc uses the expression type to determine the
variable rtype. This usually works (as mypy infers the variable type
based on the first assignment and doesn't let it change afterwards),
except in this situation.

NameExpr(y)'s type is int, but the variable should really be int | None.
Fortunately, we can use the Var node's type information which is
correct... instead of blindly assuming the first assignment's type is
right.

Fixes mypyc/mypyc#964

mypyc/irbuild/builder.py

@@ -48,6 +48,7 @@
 )
 from mypy.types import (
     AnyType,
+    DeletedType,
     Instance,
     ProperType,
     TupleType,
@@ -573,21 +574,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: