Store the type for assignment expr targets

[llchan]

Description

Assignment (walrus) expressions don't currently store their target types into the checker's type map, leading to an internal error when mypy goes to infer the type of that name downstream. This PR stores the assignment target's type explicitly. If this is not the right approach, feel free to comment/close this. I just wanted to push this up so there's a demo of a potential fix.

Fixes #9054

Test Plan

Still a work in progress, need to add tests. A test has been added that reproduces the OP's report, and this PR fixes it.

[llchan]

cc @gvanrossum since you brought this ticket to my attention. Someone should review to see if this is the right solution. At the very least it seems to fix the example presented by the OP.

[gvanrossum]

This is a good start! From the CI failure we learn that there's exactly one test that fails: testWalrusPartialTypes (at the end of test-data/unit/check-python38.test). This tests that the type of a variable initialized with an empty list is guessed if there's a following operation that appends or inserts into the list. (Same as when the initialization occurs in a regular assignment.) Apparently this code is disturbed by your small change. Maybe you can investigate a bit more? (I can give you more tips if you want to.)

[llchan]

Okay, this is green now for both the testWalrusPartialTypes test plus the new testWalrusExpr test that replicates the OP's issue.

[gvanrossum]

Thank you so much! I'm swamped but will review over the weekend (unless @hauntsaninja beats me to it :-).

[gvanrossum]

Thanks! You should be proud of yourself. Do you want me to pick another issue for you to try?

[hauntsaninja]

Thanks for fixing this! I was a little bit curious about the has_uninhabited_component... it felt a little ad hoc.

If you're willing to accept a small suggestion, I think it might be nice to fix by moving self.chk.store_type(e.target, value) over here, in visit_assignment_expr:

mypy/mypy/checkexpr.py

Line 2817 in 0b4a2c9

self.chk.check_final(e)

(visit_assignment_expr gets called by the typ = node.accept(self) on line 3782)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 0cc2f13e..2224d5e1 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -2815,6 +2815,7 @@ class ExpressionChecker(ExpressionVisitor[Type]):
         value = self.accept(e.value)
         self.chk.check_assignment(e.target, e.value)
         self.chk.check_final(e)
+        self.chk.store_type(e.target, value)
         self.find_partial_type_ref_fast_path(e.target)
         return value

It looks like find_partial_type_ref_fast_path makes its own call to self.chk.store_type, which is what we overwrote here before we added the has_uninhabited_component check, hence previous test failure.

[gvanrossum]

Sorry I didn't wait for you @hauntsaninja.

[hauntsaninja]

Not at all, thank you for being proactive!