fix second generic for AbstractContextManager being ignored when de…

…termining whether the context manager can suppress exceptions
DetachHead · Dec 21, 2024 · 10d8466 · 10d8466
1 parent 9b874de
commit 10d8466
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 19 deletions.
diff --git a/packages/pyright-internal/src/analyzer/codeFlowEngine.ts b/packages/pyright-internal/src/analyzer/codeFlowEngine.ts
@@ -1946,27 +1946,28 @@ export function getCodeFlowEngine(
                 const exitMethodName = isAsync ? '__aexit__' : '__exit__';
                 const exitType = evaluator.getBoundMagicMethod(cmType, exitMethodName);
 
-                if (exitType && isFunction(exitType) && exitType.shared.declaredReturnType) {
-                    let returnType = exitType.shared.declaredReturnType;
-
-                    // If it's an __aexit__ method, its return type will typically be wrapped
-                    // in a Coroutine, so we need to extract the return type from the third
-                    // type argument.
-                    if (isAsync) {
-                        if (
-                            isClassInstance(returnType) &&
-                            ClassType.isBuiltIn(returnType, 'Coroutine') &&
-                            returnType.priv.typeArgs &&
-                            returnType.priv.typeArgs.length >= 3
-                        ) {
-                            returnType = returnType.priv.typeArgs[2];
+                if (exitType && isFunction(exitType)) {
+                    let returnType = FunctionType.getEffectiveReturnType(exitType);
+                    if (returnType) {
+                        // If it's an __aexit__ method, its return type will typically be wrapped
+                        // in a Coroutine, so we need to extract the return type from the third
+                        // type argument.
+                        if (isAsync) {
+                            if (
+                                isClassInstance(returnType) &&
+                                ClassType.isBuiltIn(returnType, 'Coroutine') &&
+                                returnType.priv.typeArgs &&
+                                returnType.priv.typeArgs.length >= 3
+                            ) {
+                                returnType = returnType.priv.typeArgs[2];
+                            }
                         }
-                    }
 
-                    cmSwallowsExceptions = false;
-                    if (isClassInstance(returnType) && ClassType.isBuiltIn(returnType, 'bool')) {
-                        if (returnType.priv.literalValue === undefined || returnType.priv.literalValue === true) {
-                            cmSwallowsExceptions = true;
+                        cmSwallowsExceptions = false;
+                        if (isClassInstance(returnType) && ClassType.isBuiltIn(returnType, 'bool')) {
+                            if (returnType.priv.literalValue === undefined || returnType.priv.literalValue === true) {
+                                cmSwallowsExceptions = true;
+                            }
                         }
                     }
                 }

diff --git a/packages/pyright-internal/src/tests/checker.test.ts b/packages/pyright-internal/src/tests/checker.test.ts
@@ -200,6 +200,14 @@ test('With6', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('`AbstractContextManager` with exit type specified with a generic', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['withGenericExitType.py']);
+
+    TestUtils.validateResultsButBased(analysisResults, {
+        hints: [{ code: DiagnosticRule.reportUnreachable, line: 11 }],
+    });
+});
+
 test('Mro1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['mro1.py']);
 

diff --git a/packages/pyright-internal/src/tests/samples/withGenericExitType.py b/packages/pyright-internal/src/tests/samples/withGenericExitType.py
@@ -0,0 +1,12 @@
+import contextlib
+from typing import Literal
+
+def _(cm: contextlib.AbstractContextManager[object, Literal[True]]):
+    with cm:
+        raise Exception   
+    print(1)  # reachable
+
+def _(cm: contextlib.AbstractContextManager[object, Literal[False]]):
+    with cm:
+        raise Exception
+    print(1)  # error: unreachable