Use a single pass for core narrowing logic, add comments (#20659)

This PR should not change semantics

We refactor to avoid the separate loops for `type_targets` and
`value_targets`, which allows the logic to be more consolidated, helps
with future changes and is probably faster

Author: hauntsaninja

mypy/checker.py

@@ -6644,6 +6644,17 @@ def narrow_type_by_identity_equality(
         # in the context of other type checker behaviour.
         should_coerce_literals: bool
 
+        # custom_eq_indices:
+        # Operands at these indices define a custom `__eq__`. These can do arbitrary things, so we
+        # have to be more careful about what narrowing we can conclude from a successful comparison
+        custom_eq_indices: set[int]
+
+        # enum_comparison_is_ambiguous:
+        # `if x is Fruits.APPLE` we know `x` is `Fruits.APPLE`, but `if x == Fruits.APPLE: ...`
+        # it could e.g. be an int or str if Fruits is an IntEnum or StrEnum.
+        # See ambiguous_enum_equality_keys for more details
+        enum_comparison_is_ambiguous: bool
+
         if operator in {"is", "is not"}:
             is_target_for_value_narrowing = is_singleton_identity_type
             should_coerce_literals = True
@@ -6665,91 +6676,103 @@ def narrow_type_by_identity_equality(
         else:
             raise AssertionError
 
-        value_targets = []
-        type_targets = []
+        partial_type_maps = []
+
+        # For each narrowable index, we see what we can narrow based on each relevant target
         for i in expr_indices:
-            expr_type = operand_types[i]
-            if should_coerce_literals:
-                # TODO: doing this prevents narrowing a single-member Enum to literal
-                # of its member, because we expand it here and then refuse to add equal
-                # types to typemaps. As a result, `x: Foo; x == Foo.A` does not narrow
-                # `x` to `Literal[Foo.A]` iff `Foo` has exactly one member.
-                # See testMatchEnumSingleChoice
-                expr_type = coerce_to_literal(expr_type)
+            if i not in narrowable_indices:
+                continue
             if i in custom_eq_indices:
-                # We can't use types with custom __eq__ as targets for narrowing
-                # E.g. if (x: int | None) == (y: CustomEq | None), we cannot narrow x to None
+                # Handled later
                 continue
-            if is_target_for_value_narrowing(get_proper_type(expr_type)):
-                value_targets.append((i, TypeRange(expr_type, is_upper_bound=False)))
-            else:
-                type_targets.append((i, TypeRange(expr_type, is_upper_bound=False)))
-
-        partial_type_maps = []
 
-        if value_targets:
-            for i in expr_indices:
-                if i not in narrowable_indices:
+            expr_type = operand_types[i]
+            expanded_expr_type = try_expanding_sum_type_to_union(
+                coerce_to_literal(expr_type), None
+            )
+            expr_enum_keys = ambiguous_enum_equality_keys(expr_type)
+            for j in expr_indices:
+                if i == j:
                     continue
-                if i in custom_eq_indices:
-                    # Handled later
+                if j in custom_eq_indices:
+                    # We can't use types with custom __eq__ as targets for narrowing
+                    # E.g. if (x: int | None) == (y: CustomEq | None), we cannot narrow x to None
                     continue
-                expr_type = operand_types[i]
-                expr_type = coerce_to_literal(expr_type)
-                expr_type = try_expanding_sum_type_to_union(expr_type, None)
-                expr_enum_keys = ambiguous_enum_equality_keys(expr_type)
-                for j, target in value_targets:
-                    if i == j:
-                        continue
-                    if (
-                        # See comments in ambiguous_enum_equality_keys
-                        enum_comparison_is_ambiguous
-                        and len(expr_enum_keys | ambiguous_enum_equality_keys(target.item)) > 1
-                    ):
-                        continue
+                target_type = operand_types[j]
+                if should_coerce_literals:
+                    # TODO: doing this prevents narrowing a single-member Enum to literal
+                    # of its member, because we expand it here and then refuse to add equal
+                    # types to typemaps. As a result, `x: Foo; x == Foo.A` does not narrow
+                    # `x` to `Literal[Foo.A]` iff `Foo` has exactly one member.
+                    # See testMatchEnumSingleChoice
+                    target_type = coerce_to_literal(target_type)
+
+                if (
+                    # See comments in ambiguous_enum_equality_keys
+                    enum_comparison_is_ambiguous
+                    and len(expr_enum_keys | ambiguous_enum_equality_keys(target_type)) > 1
+                ):
+                    continue
+
+                target = TypeRange(target_type, is_upper_bound=False)
+                is_value_target = is_target_for_value_narrowing(get_proper_type(target_type))
+
+                if is_value_target:
                     if_map, else_map = conditional_types_to_typemaps(
-                        operands[i], *conditional_types(expr_type, [target])
+                        operands[i], *conditional_types(expanded_expr_type, [target])
                     )
                     partial_type_maps.append((if_map, else_map))
-
-        if type_targets:
-            for i in expr_indices:
-                if i not in narrowable_indices:
-                    continue
-                if i in custom_eq_indices:
-                    # Handled later
-                    continue
-                expr_type = operand_types[i]
-                for j, target in type_targets:
-                    if i == j:
-                        continue
+                else:
                     if_map, else_map = conditional_types_to_typemaps(
                         operands[i], *conditional_types(expr_type, [target])
                     )
+                    # For value targets, it is safe to narrow in the negative case.
+                    # e.g. if (x: Literal[5] | None) != (y: Literal[5]), we can narrow x to None
+                    # However, for non-value targets, we cannot do this narrowing,
+                    # and so we ignore else_map
+                    # e.g. if (x: str | None) != (y: str), we cannot narrow x to None
                     if if_map:
-                        # For type_targets, we cannot narrow in the negative case
-                        # e.g. if (x: str | None) != (y: str), we cannot narrow x to None
-                        else_map = {}
-                        partial_type_maps.append((if_map, else_map))
+                        partial_type_maps.append((if_map, {}))
 
+        # Handle narrowing for operands with custom __eq__ methods specially
+        # In most cases, we won't be able to do any narrowing
         for i in custom_eq_indices:
             if i not in narrowable_indices:
                 continue
             union_expr_type = get_proper_type(operand_types[i])
             if not isinstance(union_expr_type, UnionType):
+                # Here we won't be able to do any positive narrowing, because we can't conclude
+                # anything from a custom __eq__ returning True.
+                # But we might be able to do some negative narrowing, since we can assume
+                # a custom __eq__ is reflexive. This should only apply to custom __eq__ enums,
+                # see testNarrowingEqualityCustomEqualityEnum
                 expr_type = operand_types[i]
-                for j, target in value_targets:
-                    _if_map, else_map = conditional_types_to_typemaps(
-                        operands[i], *conditional_types(expr_type, [target])
-                    )
-                    if else_map:
-                        partial_type_maps.append(({}, else_map))
+                for j in expr_indices:
+                    if j in custom_eq_indices:
+                        continue
+                    target_type = operand_types[j]
+                    if should_coerce_literals:
+                        target_type = coerce_to_literal(target_type)
+                    target = TypeRange(target_type, is_upper_bound=False)
+                    is_value_target = is_target_for_value_narrowing(get_proper_type(target_type))
+
+                    if is_value_target:
+                        if_map, else_map = conditional_types_to_typemaps(
+                            operands[i], *conditional_types(expr_type, [target])
+                        )
+                        if else_map:
+                            partial_type_maps.append(({}, else_map))
                 continue
 
+            # If our operand with custom __eq__ is a union, where only some members of the union
+            # implement custom __eq__, then we can narrow down the other members as usual.
+            # This is basically the same logic as the main narrowing loop above.
             or_if_maps: list[TypeMap] = []
             or_else_maps: list[TypeMap] = []
             for expr_type in union_expr_type.items:
                 if has_custom_eq_checks(expr_type):
+                    # Always include union items with custom __eq__ in the type
+                    # we narrow to in the if_map
                     or_if_maps.append({operands[i]: expr_type})
 
                 for j in expr_indices:
@@ -6784,6 +6807,8 @@ def narrow_type_by_identity_equality(
 
             partial_type_maps.append((final_if_map, final_else_map))
 
+        # Handle narrowing for comparisons that produce additional narrowing, like
+        # `type(x) == T` or `x.__class__ is T`
         for i in expr_indices:
             type_expr = operands[i]
             if (