From 032f2d88cd39ee9e5a8f6fcaa2b7b86a6ea55d5d Mon Sep 17 00:00:00 2001
From: Richard Si <sichard26@gmail.com>
Date: Sun, 7 May 2023 11:20:52 -0400
Subject: [PATCH 1/2] [mypyc] Use correct rtype for variable with inferred
 optional type

```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.
---
 mypyc/irbuild/builder.py      | 12 +++++++-----
 mypyc/test-data/run-misc.test | 27 +++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index f071cc20f6b6..0b6f4ba97038 100644
--- a/mypyc/irbuild/builder.py
+++ b/mypyc/irbuild/builder.py
@@ -48,6 +48,7 @@
 )
 from mypy.types import (
     AnyType,
+    DeletedType,
     Instance,
     ProperType,
     TupleType,
@@ -570,6 +571,10 @@ 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
@@ -577,14 +582,11 @@ def get_assignment_target(
                     # 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)
diff --git a/mypyc/test-data/run-misc.test b/mypyc/test-data/run-misc.test
index 267a3441808f..8c237d278704 100644
--- a/mypyc/test-data/run-misc.test
+++ b/mypyc/test-data/run-misc.test
@@ -84,6 +84,33 @@ 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
+
+def gen(b: bool) -> Generator[Any, None, None]:
+    if b:
+        y = 1
+    else:
+        y = None
+    yield y
+
+f(False)
+next(gen(False))
+
 [case testWith]
 from typing import Any
 class Thing:

From f433dda6562cd50a2bdbe50dc47fcf0fa0114fac Mon Sep 17 00:00:00 2001
From: Richard Si <sichard26@gmail.com>
Date: Sat, 27 May 2023 11:05:30 -0400
Subject: [PATCH 2/2] Expand tests according to PR review

---
 mypyc/test-data/run-misc.test | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mypyc/test-data/run-misc.test b/mypyc/test-data/run-misc.test
index 8c237d278704..20ae76835a91 100644
--- a/mypyc/test-data/run-misc.test
+++ b/mypyc/test-data/run-misc.test
@@ -100,6 +100,7 @@ def f(b: bool) -> Any:
 
     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:
@@ -108,8 +109,10 @@ def gen(b: bool) -> Generator[Any, None, None]:
         y = None
     yield y
 
-f(False)
-next(gen(False))
+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