fix context managers that return bool | None incorrectly being treated as if they can never suppress exceptions

packages/pyright-internal/src/analyzer/codeFlowEngine.ts

@@ -55,6 +55,7 @@ import {
     isOverloadedFunction,
     isTypeSame,
     isTypeVar,
+    isUnion,
     maxTypeRecursionCount,
     NeverType,
     OverloadedFunctionType,
@@ -1777,8 +1778,15 @@ export function getCodeFlowEngine(
                     }
 
                     cmSwallowsExceptions = false;
-                    if (isClassInstance(returnType) && ClassType.isBuiltIn(returnType, 'bool')) {
-                        if (returnType.literalValue === undefined || returnType.literalValue === true) {
+                    // valid return types here are `bool | None`. if the context manager returns `True` then it suppresses,
+                    // meaning we only know for sure that the context manager can't swallow exceptions if its return type
+                    // does not allow `True`.
+                    const typesToCheck = isUnion(returnType) ? returnType.subtypes : [returnType];
+                    const boolType = typesToCheck.find(
+                        (type): type is ClassType => isClassInstance(type) && ClassType.isBuiltIn(type, 'bool')
+                    );
+                    if (boolType) {
+                        if (boolType.literalValue === undefined || boolType.literalValue === true) {
                             cmSwallowsExceptions = true;
                         }
                     }

packages/pyright-internal/src/tests/checker.test.ts

@@ -163,6 +163,14 @@ test('With2', () => {
     TestUtils.validateResults(analysisResults, 3);
 });
 
+test('context manager where __exit__ returns bool | None', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['withBased.py']);
+    TestUtils.validateResultsButBased(analysisResults, {
+        unreachableCodes: [{ line: 47 }, { line: 62 }],
+        unusedCodes: undefined,
+    });
+});
+
 test('With3', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['with3.py']);
 

packages/pyright-internal/src/tests/samples/withBased.py

@@ -0,0 +1,63 @@
+import contextlib
+from types import TracebackType
+from typing import Iterator, Literal
+
+from typing_extensions import assert_never
+
+class BoolOrNone(contextlib.AbstractContextManager[None]):
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> bool | None:
+        ...
+
+def _():
+    with BoolOrNone():
+        raise Exception
+    print(1)  # reachable
+
+class TrueOrNone(contextlib.AbstractContextManager[None]):
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> Literal[True] | None:
+        ...
+
+def _():
+    with TrueOrNone():
+        raise Exception
+    print(1)  # reachable
+
+
+class FalseOrNone(contextlib.AbstractContextManager[None]):
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> Literal[False] | None:
+        ...
+
+def _():
+    with FalseOrNone():
+        raise Exception
+    print(1)  # unreachable
+
+
+class OnlyNone(contextlib.AbstractContextManager[None]):
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> None:
+        ...
+
+def _():
+    with OnlyNone():
+        raise Exception
+    print(1)  # unreachable

packages/pyright-internal/src/tests/testUtils.ts

@@ -245,7 +245,11 @@ export const validateResultsButBased = (allResults: FileAnalysisResult[], expect
                 code: result.getRule() as DiagnosticRule | undefined,
             })
         );
-        const expectedResult = expectedResults[diagnosticType] ?? [];
-        expect(new Set(actualResult)).toEqual(new Set(expectedResult.map(expect.objectContaining)));
+        const expectedResult = expectedResults[diagnosticType];
+        // if it's explicitly in the expected results as undefined, that means we don't care.
+        // if it's not in the expected results at all, then check it
+        if (!(diagnosticType in expectedResults) || expectedResult !== undefined) {
+            expect(new Set(actualResult)).toEqual(new Set((expectedResult ?? []).map(expect.objectContaining)));
+        }
     }
 };