[ty] improve complex TDD-based narrowing #23201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

mtshiba wants to merge 39 commits into astral-sh:main from mtshiba:improve-23109

+1,078 −336

crates/ty_python_semantic/resources/mdtest/binary/integers.md

-Original file line number
+Diff line change
@@ -1,5 +1,11 @@
     # Binary operations on integers
+    > Developer's note: This is mainly a test for the behavior of the type inferer. The constant
+    > evaluator (`resolve_to_literal`) of `SemanticIndexBuilder` is implemented separately from the type
+    > inferer, so if you modify the contents of this file or the type inferer, please also modify the
+    > implementation of `resolve_to_literal` and the unit tests (semantic_index/tests/const_eval\_\*) at
+    > the same time.
     ## Basic Arithmetic
     ```py
@@ Expand Down @@

crates/ty_python_semantic/resources/mdtest/loops/while_loop.md

-Original file line number
+Diff line change
@@ Expand Up @@
     def random() -> bool:
         return True
-    x = random()
-    reveal_type(x)  # revealed: bool
-    while x:
-        pass
-    reveal_type(x)  # revealed: Literal[False]
+    def _(x: bool):
+        while x:
+            pass
+        reveal_type(x)  # revealed: Literal[False]
     ```
     However, we can't narrow `x` like this when there's a `break` in the loop:
     ```py
-    x = random()
-    while x:
-        if random():
+    def _(x: bool):
+        while x:
+            if random():
+                break
+        reveal_type(x)  # revealed: bool
+    def _(x: bool):
+        while x:
+            pass
+        reveal_type(x)  # revealed: Literal[False]
+        x = random()
+        while x:
+            if random():
+                break
+        reveal_type(x)  # revealed: bool
+    def _(y: int | None):
+        x = 1
+        while True:
+            if x == 0:
+                break
+            if y is None:
+                y = 0
+                continue
             break
-    reveal_type(x)  # revealed: bool
+        reveal_type(y)  # revealed: int
     ```
     ### Non-static loop conditions
@@ Expand Down @@

crates/ty_python_semantic/resources/mdtest/narrow/post_if_statement.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -180,14 +180,26 @@ def _(x: int | None): @@
     ```
     ```py
+    from typing import Final
     def _(x: int | None):
         if 1 + 1 == 2:
             if x is None:
                 return
             reveal_type(x)  # revealed: int
-        # TODO: should be `int` (the else-branch of `1 + 1 == 2` is unreachable)
-        reveal_type(x)  # revealed: int | None
+        reveal_type(x)  # revealed: int
+    # non-constant but always-true condition
+    needs_inference: Final = True
+    def _(x: int | None):
+        if needs_inference:
+            if x is None:
+                return
+            reveal_type(x)  # revealed: int
+        reveal_type(x)  # revealed: int
     ```
     This also works when the always-true condition is nested inside a narrowing branch:
@@ Expand All / @@ -198,9 +210,14 @@ def _(x: int | None): @@
             if 1 + 1 == 2:
                 return
-        # TODO: should be `int` (the inner always-true branch makes the outer
-        # if-branch terminal)
-        reveal_type(x)  # revealed: int | None
+        reveal_type(x)  # revealed: int
+    def _(x: int | None):
+        if x is None:
+            if needs_inference:
+                return
+        reveal_type(x)  # revealed: int
     ```
     ## Narrowing from `assert` should not affect reassigned variables
@@ Expand Down @@

crates/ty_python_semantic/src/semantic_index.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -925,6 +925,50 @@ mod tests { @@
                 .collect()
         }
+        /// A function to test how the constant evaluator of `SemanticIndexBuilder` evaluates an expression
+        /// (the evaluation should match that of `TypeInferenceBuilder`).
+        /// For example, for the input `x = 1\nif cond: x = 2\nx`, if `cond` evaluates to `AlwaysTrue`, it returns `vec![2]`,
+        /// if it evaluates to `AlwaysFalse`, it returns `vec![1]`, if it evaluates to `Ambiguous`, it returns `vec![1, 2]`.
+        fn reachable_bindings_for_terminal_use(content: &str) -> Vec<i64> {
+            let TestCase { db, file } = test_case(content);
+            let scope = global_scope(&db, file);
+            let module = parsed_module(&db, file).load(&db);
+            let ast = module.syntax();
+            let terminal_expr = ast
+                .body
+                .last()
+                .and_then(ast::Stmt::as_expr_stmt)
+                .map(|stmt| stmt.value.as_ref())
+                .expect("expected terminal expression statement");
+            let terminal_name = terminal_expr
+                .as_name_expr()
+                .expect("terminal expression should be a name");
+            let use_id = terminal_name.scoped_use_id(&db, scope);
+            let use_def = use_def_map(&db, scope);
+            use_def
+                .bindings_at_use(use_id)
+                .filter_map(|binding_with_constraints| {
+                    let definition = binding_with_constraints.binding.definition()?;
+                    let DefinitionKind::Assignment(assignment) = definition.kind(&db) else {
+                        return None;
+                    };
+                    let ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
+                        value: ast::Number::Int(value),
+                        ..
+                    }) = assignment.value(&module)
+                    else {
+                        return None;
+                    };
+                    value.as_i64()
+                })
+                .collect::<Vec<_>>()
+        }
         #[test]
         fn empty() {
             let TestCase { db, file } = test_case("");
@@ Expand Down Expand Up / @@ -1590,6 +1634,71 @@ class C[T]: @@
             assert_eq!(*num, 1);
         }
+        #[test]
+        fn const_eval_lshift_overflow_is_ambiguous() {
+            let values = reachable_bindings_for_terminal_use(
+                "
+    x = 1
+    if 1 << 63:
+        x = 2
+    x
+    ",
+            );
+            assert_eq!(values, vec![1, 2]);
+        }
+        #[test]
+        fn const_eval_lshift_zero_short_circuit() {
+            let values = reachable_bindings_for_terminal_use(
+                "
+    x = 1
+    if 0 << 4000000000000000000:
+        x = 2
+    x
+    ",
+            );
+            assert_eq!(values, vec![1]);
+        }
+        #[test]
+        fn const_eval_rshift_large_positive() {
+            let values = reachable_bindings_for_terminal_use(
+                "
+    x = 1
+    if 1 >> 5000000000:
+        x = 2
+    x
+    ",
+            );
+            assert_eq!(values, vec![1]);
+        }
+        #[test]
+        fn const_eval_rshift_large_negative_operand() {
+            let values = reachable_bindings_for_terminal_use(
+                "
+    x = 1
+    if (-1) >> 5000000000:
+        x = 2
+    x
+    ",
+            );
+            assert_eq!(values, vec![2]);
+        }
+        #[test]
+        fn const_eval_negative_lshift_is_ambiguous() {
+            let values = reachable_bindings_for_terminal_use(
+                "
+    x = 1
+    if 42 << -3:
+        x = 2
+    x
+    ",
+            );
+            assert_eq!(values, vec![1, 2]);
+        }
         #[test]
         fn expression_scope() {
             let TestCase { db, file } = test_case("x = 1;\ndef test():\n  y = 4");
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ty] improve complex TDD-based narrowing #23201

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

[ty] improve complex TDD-based narrowing #23201

Are you sure you want to change the base?

[ty] improve complex TDD-based narrowing #23201

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!