Fix internal error on list/dict comprehension with walrus operator in global scope

[dhood]

Fixes #8986

The reported issue is for list comprehensions involving assignment expressions, though I noticed that dictionary comprehensions are also impacted.

# mwe.py
lst = [1, 2, 3, 4]
processed_list = [e for elt in lst if (e := elt - 1) > 1]

dct = {'a': 1, 'b': 2, 'c': 3, 'd': 4}
processed_dict = {k: new_v for k, v in dct.items() if (new_v := v - 1) > 1

On the other hand, the same declarations inside a function pass:

# valid_comprehensions.py
def calling_function() -> None:
    lst = [1, 2, 3, 4]
    processed_list = [e for elt in lst if (e := elt - 1) > 1]
    reveal_type(processed_list)

    dct = {'a': 1, 'b': 2, 'c': 3, 'd': 4}
    processed_dict = {k: new_v for k, v in dct.items() if (new_v := v - 1) > 1}
    reveal_type(processed_dict)

This tipped me off to it being related to the special logic added in #6899 to store the variables in the scope of the comprehension's caller, which I believe can benefit from special treatment for the edge case of the caller being in the global scope. That's the approach that I took, though it's my first contribution to MyPy so another approach may be more appropriate.

I added regression tests for these cases, but I also added some tests that were already passing. I've kept them in separate commits (with CI triggered separately) to be clear about which ones are being addressed by this PR!

[dhood]

If reviewers are wondering why I didn't just use d.items() here instead of the list of tuples, it's because it's causing error: "Dict[str, int]" has no attribute "items".
I couldn't figure out why that was happening, I tried different flags and casting and whatnot but something in the test environment must be different... in a regular python script filtered_dict = {k: new_v for k, v in d.items() if (new_v := v + 1) == 2} can have its type revealed without such an error. Open to suggestions if you have ideas what's causing that!

[JelleZijlstra]

This is almost certainly because the test fixture you're using happens not to define dict.items. What you're doing here is fine.

[JelleZijlstra]

Thanks for catching this! Sorry for missing it when I wrote the code. Note that the specification for this behavior is in https://www.python.org/dev/peps/pep-0572/#scope-of-the-target.

Could you also add some tests with assignments happening in class scopes?

[JelleZijlstra]

Why is this cast() needed? Could you use a # type: comment on the first assignment instead?

[dhood]

the cast here was because self.locals is a list of Optional[SymbolTable], and so without it, you would get:

mypy\semanal.py:4552: error: Incompatible types in assignment (expression has type "Optional[SymbolTable]", variable has type "SymbolTable")  [assignment]
                                names = self.locals[-1 - i]
                                        ^

I reasoned (at the time) that it is a safe cast as long as we are not retrieving the None that is the first element in the list.

mypy/mypy/semanal.py

Line 226 in ed9ba4f

self.locals = [None]

However, now that you have mentioned to check the class scope, I have realised that there may be other elements in the list that are None:

mypy/mypy/semanal.py

Line 1166 in ed9ba4f

self.locals.append(None) # Add class scope

Perhaps once I deal with the class case properly this cast will not be necessary. Either way, for now, my original assumption that only the first entry is None seems invalid! So I will find an alternative approach.

[dhood]

Since comprehensions in a class scope e.g.

class Example:
    [(j := i) for i in range(5)]

can also result in the retrieved self.locals[-1 -i] symbol table being None, I added extra logic for that case, which removed the need for the cast, in a68bd68.

Because this specific case is called out in https://www.python.org/dev/peps/pep-0572/#scope-of-the-target as being invalid syntax (I know that you know this, I am just being verbose), the specific logic that's added is actually just an assertion that we don't end up in that case.

[dhood]

Sorry for missing it when I wrote the code.

No, not at all. Certainly no anger or frustration from me, I am grateful of the initial implementation that you wrote! (which worked great in most cases, this is just an edge case). Actually I was excited to stumble upon this bug... I am very happy to have something small to contribute towards the support of the walrus operator! 😄

Will add those class-scope tests and ping you again.

[dhood]

@JelleZijlstra, I added tests for the class scope in 8585a4f

In doing so, I noticed some cases of what I consider inappropriate mypy behaviour:

1. The invalid syntax highlighted in https://www.python.org/dev/peps/pep-0572/#scope-of-the-target where a comprehension is used in a class scope does not raise a Syntax error. I opened PEP 572: Invalid python syntax for walrus operator not raising parser error #9095
2. Use of assignment operator within a list and subsequent re-use of the variable in the same list e.g. l = [x2 := 1, x2 + 2] from a class scope does not work. This is another variation of Walrus operators unrecognised in dictionary initialisation of a class attribute #8379, so I have commented there.
3. Re-use of variables defined in comprehensions (in any scope) is not permitted. Opened PEP 572: re-use of variable defined with assignment expression / walrus operator not permitted #9096

Since I consider these to be separate issues (they happen on master also), and because I didn't plan time to address those this weekend, I have spun them off to not block this PR.

[dhood]

thanks for the review, @JelleZijlstra - I'll try to make progress on the spin-off issues this weekend