[partially defined] fix a false-negative with variable defined in a skipped branch (#14221)

The goal of partially-defined check is to detect variables that could be
undefined but semantic analyzer doesn't detect them as undefined. In
this case, a variable was defined in a branch that returns, so semantic
analyzer considered it defined when it was not.

I've discovered this when testing support for try/except statements
(#14114).

Author: ilinum

mypy/partially_defined.py

@@ -121,23 +121,29 @@ def is_defined_in_a_branch(self, name: str) -> bool:
         return False
 
     def done(self) -> BranchState:
-        branches = [b for b in self.branches if not b.skipped]
-        if len(branches) == 0:
-            return BranchState(skipped=True)
-        if len(branches) == 1:
-            return branches[0]
-
-        # must_be_defined is a union of must_be_defined of all branches.
-        must_be_defined = set(branches[0].must_be_defined)
-        for b in branches[1:]:
-            must_be_defined.intersection_update(b.must_be_defined)
-        # may_be_defined are all variables that are not must be defined.
+        # First, compute all vars, including skipped branches. We include skipped branches
+        # because our goal is to capture all variables that semantic analyzer would
+        # consider defined.
         all_vars = set()
-        for b in branches:
+        for b in self.branches:
             all_vars.update(b.may_be_defined)
             all_vars.update(b.must_be_defined)
+        # For the rest of the things, we only care about branches that weren't skipped.
+        non_skipped_branches = [b for b in self.branches if not b.skipped]
+        if len(non_skipped_branches) > 0:
+            must_be_defined = non_skipped_branches[0].must_be_defined
+            for b in non_skipped_branches[1:]:
+                must_be_defined.intersection_update(b.must_be_defined)
+        else:
+            must_be_defined = set()
+        # Everything that wasn't defined in all branches but was defined
+        # in at least one branch should be in `may_be_defined`!
         may_be_defined = all_vars.difference(must_be_defined)
-        return BranchState(may_be_defined=may_be_defined, must_be_defined=must_be_defined)
+        return BranchState(
+            must_be_defined=must_be_defined,
+            may_be_defined=may_be_defined,
+            skipped=len(non_skipped_branches) == 0,
+        )
 
 
 class Scope:

test-data/unit/check-partially-defined.test

@@ -345,6 +345,12 @@ def f5() -> int:
         return 3
     return 1
 
+def f6() -> int:
+    if int():
+        x = 0
+        return x
+    return x  # E: Name "x" may be undefined
+
 [case testDefinedDifferentBranchUseBeforeDef]
 # flags: --enable-error-code partially-defined --enable-error-code use-before-def