Added type narrowing support for index expressions where the index value is a string literal.

docs/type-concepts.md

@@ -94,7 +94,7 @@ For more details about generic types, type parameters, and invariance, refer to
 
 ### Type Narrowing
 
-Pyright uses a technique called “type narrowing” to track the type of a symbol based on code flow. Consider the following code:
+Pyright uses a technique called “type narrowing” to track the type of an expression based on code flow. Consider the following code:
 
 ```python
 val_str: str = "hi"
@@ -126,6 +126,23 @@ Another assignment occurs several lines further down, this time within a conditi
 
 Another way that types can be narrowed is through the use of conditional code flow statements like `if`, `while`, and `assert`. Type narrowing applies to the block of code that is “guarded” by that condition, so type narrowing in this context is sometimes referred to as a “type guard”. For example, if you see the conditional statement `if x is None:`, the code within that `if` statement can assume that `x` contains `None`. Within the code sample above, we see an example of a type guard involving a call to `isinstance`. The type checker knows that `isinstance(val, int)` will return True only in the case where `val` contains a value of type `int`, not type `str`. So the code within the `if` block can assume that `val` contains a value of type `int`, and the code within the `else` block can assume that `val` contains a value of type `str`. This demonstrates how a type (in this case `Union[int, str]`) can be narrowed in both a positive (`if`) and negative (`else`) test.
 
+The following expression forms support type narrowing:
+
+* `<ident>` (where `<ident>` is an identifier)
+* `<expr>.<member>` (member access expression where `<expr>` is a supported expression form)
+* `<expr>[<int>]` (subscript expression where `<int>` is a non-negative integer)
+* `<expr>[<str>]` (subscript expression where `<str>` is a string literal)
+
+Examples of expressions that support type narrowing:
+
+* `my_var`
+* `employee.name`
+* `a.foo.next`
+* `args[3]`
+* `kwargs["bar"]`
+* `a.b.c[3]["x"].d`
+
+
 ### Type Guards
 
 In addition to assignment-based type narrowing, Pyright supports the following type guards.

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

@@ -13,7 +13,7 @@
  * TypeScript compiler.
  */
 
-import { assert } from '../common/debug';
+import { assert, fail } from '../common/debug';
 import {
     ArgumentCategory,
     CallNode,
@@ -25,6 +25,7 @@ import {
     NameNode,
     NumberNode,
     ParseNodeType,
+    StringNode,
     SuiteNode,
 } from '../parser/parseNodes';
 
@@ -156,7 +157,7 @@ export function isCodeFlowSupportedForReference(reference: ExpressionNode): bool
 
     if (reference.nodeType === ParseNodeType.Index) {
         // Allow index expressions that have a single subscript that is a
-        // literal integer value.
+        // literal integer or string value.
         if (
             reference.items.length !== 1 ||
             reference.trailingComma ||
@@ -167,7 +168,14 @@ export function isCodeFlowSupportedForReference(reference: ExpressionNode): bool
         }
 
         const subscriptNode = reference.items[0].valueExpression;
-        if (subscriptNode.nodeType !== ParseNodeType.Number || subscriptNode.isImaginary || !subscriptNode.isInteger) {
+        const isIntegerIndex =
+            subscriptNode.nodeType === ParseNodeType.Number && !subscriptNode.isImaginary && subscriptNode.isInteger;
+        const isStringIndex =
+            subscriptNode.nodeType === ParseNodeType.StringList &&
+            subscriptNode.strings.length === 1 &&
+            subscriptNode.strings[0].nodeType === ParseNodeType.String;
+
+        if (!isIntegerIndex && !isStringIndex) {
             return false;
         }
 
@@ -184,10 +192,20 @@ export function createKeyForReference(reference: CodeFlowReferenceExpressionNode
     } else if (reference.nodeType === ParseNodeType.MemberAccess) {
         const leftKey = createKeyForReference(reference.leftExpression as CodeFlowReferenceExpressionNode);
         key = `${leftKey}.${reference.memberName.value}`;
-    } else {
+    } else if (reference.nodeType === ParseNodeType.Index) {
         const leftKey = createKeyForReference(reference.baseExpression as CodeFlowReferenceExpressionNode);
-        assert(reference.items.length === 1 && reference.items[0].valueExpression.nodeType === ParseNodeType.Number);
-        key = `${leftKey}[${(reference.items[0].valueExpression as NumberNode).value.toString()}]`;
+        assert(reference.items.length === 1);
+        if (reference.items[0].valueExpression.nodeType === ParseNodeType.Number) {
+            key = `${leftKey}[${(reference.items[0].valueExpression as NumberNode).value.toString()}]`;
+        } else if (reference.items[0].valueExpression.nodeType === ParseNodeType.StringList) {
+            const valExpr = reference.items[0].valueExpression;
+            assert(valExpr.strings.length === 1 && valExpr.strings[0].nodeType === ParseNodeType.String);
+            key = `${leftKey}["${(valExpr.strings[0] as StringNode).value}"]`;
+        } else {
+            fail('createKeyForReference received unexpected index type');
+        }
+    } else {
+        fail('createKeyForReference received unexpected expression type');
     }
 
     return key;

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

@@ -27,7 +27,6 @@ import {
     LambdaNode,
     ModuleNode,
     NameNode,
-    NumberNode,
     ParameterCategory,
     ParseNode,
     ParseNodeType,
@@ -859,13 +858,35 @@ export function isMatchingExpression(reference: ExpressionNode, expression: Expr
             return false;
         }
 
-        const referenceNumberNode = reference.items[0].valueExpression as NumberNode;
-        const subscriptNode = expression.items[0].valueExpression;
-        if (subscriptNode.nodeType !== ParseNodeType.Number || subscriptNode.isImaginary || !subscriptNode.isInteger) {
-            return false;
+        if (reference.items[0].valueExpression.nodeType === ParseNodeType.Number) {
+            const referenceNumberNode = reference.items[0].valueExpression;
+            const subscriptNode = expression.items[0].valueExpression;
+            if (
+                subscriptNode.nodeType !== ParseNodeType.Number ||
+                subscriptNode.isImaginary ||
+                !subscriptNode.isInteger
+            ) {
+                return false;
+            }
+
+            return referenceNumberNode.value === subscriptNode.value;
         }
 
-        return referenceNumberNode.value === subscriptNode.value;
+        if (reference.items[0].valueExpression.nodeType === ParseNodeType.StringList) {
+            const referenceStringListNode = reference.items[0].valueExpression;
+            const subscriptNode = expression.items[0].valueExpression;
+            if (
+                referenceStringListNode.strings.length === 1 &&
+                referenceStringListNode.strings[0].nodeType === ParseNodeType.String &&
+                subscriptNode.nodeType === ParseNodeType.StringList &&
+                subscriptNode.strings.length === 1 &&
+                subscriptNode.strings[0].nodeType === ParseNodeType.String
+            ) {
+                return referenceStringListNode.strings[0].value === subscriptNode.strings[0].value;
+            }
+        }
+
+        return false;
     }
 
     return false;

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

@@ -41,3 +41,25 @@ def func2(v1: List[Union[Dict[str, str], List[str]]]):
     if isinstance(v1[0], dict):
         t_v1_0: Literal["Dict[str, str]"] = reveal_type(v1[0])
         t_v1_1: Literal["Dict[str, str] | List[str]"] = reveal_type(v1[1])
+
+
+def func3():
+    v1: Dict[str, int] = {}
+
+    t_v1_0: Literal["int"] = reveal_type(v1["x1"])
+    v1["x1"] = 3
+    t_v1_1: Literal["Literal[3]"] = reveal_type(v1["x1"])
+
+    v1[f"x2"] = 5
+    t_v1_2: Literal["int"] = reveal_type(v1["x2"])
+
+    v1 = {}
+    t_v1_3: Literal["int"] = reveal_type(v1["x1"])
+
+    v2: Dict[str, Dict[str, int]] = {}
+
+    t_v2_0: Literal["int"] = reveal_type(v2["y1"]["y2"])
+    v2["y1"]["y2"] = 3
+    t_v2_1: Literal["Literal[3]"] = reveal_type(v2["y1"]["y2"])
+    v2["y1"] = {}
+    t_v2_2: Literal["int"] = reveal_type(v2["y1"]["y2"])