From 606e12e8155c5dc8ea514b5e9c7bdfd03f468ff9 Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Tue, 1 Oct 2024 13:55:47 -0700
Subject: [PATCH] Fixed bug that caused type alias to `Never` or `NoReturn` to
 lose its type alias association in some circumstances. (#9127)

---
 .../pyright-internal/src/analyzer/types.ts    | 21 ++++++++++---------
 .../src/common/collectionUtils.ts             | 16 ++++++++++++++
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/packages/pyright-internal/src/analyzer/types.ts b/packages/pyright-internal/src/analyzer/types.ts
index cc0e4cd76e7a..063e9840a116 100644
--- a/packages/pyright-internal/src/analyzer/types.ts
+++ b/packages/pyright-internal/src/analyzer/types.ts
@@ -7,6 +7,7 @@
  * Representation of types used during type analysis within Python.
  */
 
+import { partition } from '../common/collectionUtils';
 import { assert } from '../common/debug';
 import { Uri } from '../common/uri/uri';
 import { ArgumentNode, ExpressionNode, NameNode, ParamCategory } from '../parser/parseNodes';
@@ -3650,19 +3651,19 @@ export interface CombineTypesOptions {
 // are combined into a UnionType. NeverTypes are filtered out.
 // If no types remain in the end, a NeverType is returned.
 export function combineTypes(subtypes: Type[], options?: CombineTypesOptions): Type {
-    // Filter out any "Never" and "NoReturn" types.
-    let sawNoReturn = false;
+    let neverTypes: NeverType[];
 
-    if (subtypes.some((subtype) => subtype.category === TypeCategory.Never))
-        subtypes = subtypes.filter((subtype) => {
-            if (subtype.category === TypeCategory.Never && subtype.priv.isNoReturn) {
-                sawNoReturn = true;
-            }
-            return subtype.category !== TypeCategory.Never;
-        });
+    // Filter out any Never or NoReturn types.
+    [neverTypes, subtypes] = partition<Type, NeverType>(subtypes, isNever);
 
     if (subtypes.length === 0) {
-        return sawNoReturn ? NeverType.createNoReturn() : NeverType.createNever();
+        if (neverTypes.length > 0) {
+            // Prefer NoReturn over Never. This approach preserves type alias
+            // information if present.
+            return neverTypes.find((t) => t.priv.isNoReturn) ?? neverTypes[0];
+        }
+
+        return NeverType.createNever();
     }
 
     // Handle the common case where there is only one type.
diff --git a/packages/pyright-internal/src/common/collectionUtils.ts b/packages/pyright-internal/src/common/collectionUtils.ts
index 3311a09fa616..0135e0e462e0 100644
--- a/packages/pyright-internal/src/common/collectionUtils.ts
+++ b/packages/pyright-internal/src/common/collectionUtils.ts
@@ -73,6 +73,22 @@ export function appendArray<T>(to: T[], elementsToPush: T[]) {
     }
 }
 
+/** Works like Array.filter except that it returns a second array with the filtered elements. **/
+export function partition<T, S extends T>(array: readonly T[], cb: (value: T) => boolean): [S[], T[]] {
+    const trueItems: S[] = [];
+    const falseItems: T[] = [];
+
+    for (const item of array) {
+        if (cb(item)) {
+            trueItems.push(item as S);
+        } else {
+            falseItems.push(item);
+        }
+    }
+
+    return [trueItems, falseItems];
+}
+
 /** Works like Array.prototype.find, returning `undefined` if no element satisfying the predicate is found. */
 export function find<T, U extends T>(
     array: readonly T[],