Skip to content

Commit f8f9453

Browse files
authored
[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
1 parent c4f55ae commit f8f9453

File tree

2 files changed

+37
-5
lines changed

2 files changed

+37
-5
lines changed

mypyc/irbuild/builder.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
)
4949
from mypy.types import (
5050
AnyType,
51+
DeletedType,
5152
Instance,
5253
ProperType,
5354
TupleType,
@@ -573,21 +574,22 @@ def get_assignment_target(
573574
self.error("Cannot assign to the first argument of classmethod", line)
574575
if lvalue.kind == LDEF:
575576
if symbol not in self.symtables[-1]:
577+
if isinstance(symbol, Var) and not isinstance(symbol.type, DeletedType):
578+
reg_type = self.type_to_rtype(symbol.type)
579+
else:
580+
reg_type = self.node_type(lvalue)
576581
# If the function is a generator function, then first define a new variable
577582
# in the current function's environment class. Next, define a target that
578583
# refers to the newly defined variable in that environment class. Add the
579584
# target to the table containing class environment variables, as well as the
580585
# current environment.
581586
if self.fn_info.is_generator:
582587
return self.add_var_to_env_class(
583-
symbol,
584-
self.node_type(lvalue),
585-
self.fn_info.generator_class,
586-
reassign=False,
588+
symbol, reg_type, self.fn_info.generator_class, reassign=False
587589
)
588590

589591
# Otherwise define a new local variable.
590-
return self.add_local_reg(symbol, self.node_type(lvalue))
592+
return self.add_local_reg(symbol, reg_type)
591593
else:
592594
# Assign to a previously defined variable.
593595
return self.lookup(symbol)

mypyc/test-data/run-misc.test

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,36 @@ assert f(a) is a
8484
assert g(None) == 1
8585
assert g(a) == 2
8686

87+
[case testInferredOptionalAssignment]
88+
from typing import Any, Generator
89+
90+
def f(b: bool) -> Any:
91+
if b:
92+
x = None
93+
else:
94+
x = 1
95+
96+
if b:
97+
y = 1
98+
else:
99+
y = None
100+
101+
m = 1 if b else None
102+
n = None if b else 1
103+
return ((x, y), (m, n))
104+
105+
def gen(b: bool) -> Generator[Any, None, None]:
106+
if b:
107+
y = 1
108+
else:
109+
y = None
110+
yield y
111+
112+
assert f(False) == ((1, None), (None, 1))
113+
assert f(True) == ((None, 1), (1, None))
114+
assert next(gen(False)) is None
115+
assert next(gen(True)) == 1
116+
87117
[case testWith]
88118
from typing import Any
89119
class Thing:

0 commit comments

Comments
 (0)