fix Support constrained TypeVars #1422

conformance/third_party/conformance.exp

@@ -5037,17 +5037,6 @@
       "stop_column": 17,
       "stop_line": 34
     },
-    {
-      "code": -2,
-      "column": 12,
-      "concise_description": "Returned type `bytes | Unknown` is not assignable to declared return type `AnyStr`",
-      "description": "Returned type `bytes | Unknown` is not assignable to declared return type `AnyStr`",
-      "line": 34,
-      "name": "bad-return",
-      "severity": "error",
-      "stop_column": 17,
-      "stop_line": 34
-    },
     {
       "code": -2,
       "column": 15,

conformance/third_party/conformance.result

@@ -134,7 +134,7 @@
     "Line 98: Expected 1 errors"
   ],
   "generics_basic.py": [
-    "Line 34: Unexpected errors ['`+` is not supported between `AnyStr` and `AnyStr`\\n  Argument `AnyStr` is not assignable to parameter `value` with type `Buffer` in function `bytes.__add__`\\n  Protocol `Buffer` requires attribute `__buffer__`', '`+` is not supported between `AnyStr` and `AnyStr`\\n  No matching overload found for function `str.__add__` called with arguments: (AnyStr)\\n  Possible overloads:\\n  (value: LiteralString, /) -> LiteralString\\n  (value: str, /) -> str [closest match]', 'Returned type `bytes | Unknown` is not assignable to declared return type `AnyStr`']",
+    "Line 34: Unexpected errors ['`+` is not supported between `AnyStr` and `AnyStr`\\n  Argument `AnyStr` is not assignable to parameter `value` with type `Buffer` in function `bytes.__add__`\\n  Protocol `Buffer` requires attribute `__buffer__`', '`+` is not supported between `AnyStr` and `AnyStr`\\n  No matching overload found for function `str.__add__` called with arguments: (AnyStr)\\n  Possible overloads:\\n  (value: LiteralString, /) -> LiteralString\\n  (value: str, /) -> str [closest match]']",
     "Line 67: Unexpected errors ['assert_type(MyStr, str) failed']",
     "Line 68: Unexpected errors ['assert_type(MyStr, str) failed', 'Argument `str` is not assignable to parameter `y` with type `MyStr` in function `concat`']"
   ],

pyrefly/lib/binding/function.rs

@@ -22,6 +22,7 @@ use ruff_python_ast::ExceptHandler;
 use ruff_python_ast::Expr;
 use ruff_python_ast::ExprCall;
 use ruff_python_ast::Identifier;
+use ruff_python_ast::Parameter;
 use ruff_python_ast::Parameters;
 use ruff_python_ast::Stmt;
 use ruff_python_ast::StmtExpr;
@@ -32,6 +33,7 @@ use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 use starlark_map::small_map::SmallMap;
+use starlark_map::small_set::SmallSet;
 
 use crate::binding::binding::AnnotationTarget;
 use crate::binding::binding::Binding;
@@ -95,6 +97,9 @@ fn is_annotated<T>(returns: &Option<T>, params: &Parameters) -> bool {
     }
     false
 }
+
+type ConstrainedTypeVarParams = SmallMap<Name, SmallSet<Name>>;
+
 struct SelfAttrNames<'a> {
     self_name: &'a Name,
     names: SmallMap<Name, TextRange>,
@@ -350,8 +355,23 @@ impl<'a> BindingsBuilder<'a> {
     /// This function must not be called unless the function body statements will be bound;
     /// it relies on that binding to ensure we don't have a dangling `Idx<Key>` (which could lead
     /// to a panic).
-    fn implicit_return(&mut self, body: &[Stmt], func_name: &Identifier) -> Idx<Key> {
-        let last_exprs = function_last_expressions(body, self.sys_info).map(|x| {
+    fn implicit_return(
+        &mut self,
+        body: &[Stmt],
+        func_name: &Identifier,
+        parameters: &Parameters,
+    ) -> Idx<Key> {
+        let constrained_typevars = self.constrained_typevar_params(parameters);
+        let last_exprs = function_last_expressions(
+            body,
+            self.sys_info,
+            if constrained_typevars.is_empty() {
+                None
+            } else {
+                Some(&constrained_typevars)
+            },
+        )
+        .map(|x| {
             x.into_map(|(last, x)| (last, self.last_statement_idx_for_implicit_return(last, x)))
                 .into_boxed_slice()
         });
@@ -361,6 +381,67 @@ impl<'a> BindingsBuilder<'a> {
         )
     }
 
+    fn constrained_typevar_params(&self, parameters: &Parameters) -> ConstrainedTypeVarParams {
+        let mut constrained = ConstrainedTypeVarParams::new();
+        let mut consider_param = |param: &Parameter| {
+            let Some(annotation) = param.annotation() else {
+                return;
+            };
+            let Expr::Name(typevar_name) = annotation else {
+                return;
+            };
+            let Some(constraints) = self.typevar_constraints_for_name(&typevar_name.id) else {
+                return;
+            };
+            constrained.insert(param.name().id.clone(), constraints);
+        };
+        for param in parameters.iter_non_variadic_params() {
+            consider_param(&param.parameter);
+        }
+        if let Some(vararg) = &parameters.vararg {
+            consider_param(vararg);
+        }
+        if let Some(kwarg) = &parameters.kwarg {
+            consider_param(kwarg);
+        }
+        constrained
+    }
+
+    fn typevar_constraints_for_name(&self, name: &Name) -> Option<SmallSet<Name>> {
+        let (idx, _) = self.scopes.binding_idx_for_name(name)?;
+        let (_, binding) = self.get_original_binding(idx)?;
+        let binding = binding?;
+        let Binding::TypeVar(_, _, call) = binding else {
+            return None;
+        };
+        Self::typevar_constraints_from_call(call)
+    }
+
+    fn typevar_constraints_from_call(call: &ExprCall) -> Option<SmallSet<Name>> {
+        if call
+            .arguments
+            .keywords
+            .iter()
+            .any(|kw| kw.arg.as_ref().is_some_and(|id| id.id == "bound"))
+        {
+            return None;
+        }
+        let mut args = call.arguments.args.iter();
+        let _name_arg = args.next()?;
+        let mut constraints = SmallSet::new();
+        for arg in args {
+            let Expr::Name(name) = arg else {
+                return None;
+            };
+            constraints.insert(name.id.clone());
+        }
+        if constraints.len() < 2 {
+            None
+        } else {
+            Some(constraints)
+        }
+    }
+
     /// Handles both checking yield / return expressions and binding the return type.
     fn analyze_return_type(
         &mut self,
@@ -625,7 +706,7 @@ impl<'a> BindingsBuilder<'a> {
                     self_assignments
                 }
                 UntypedDefBehavior::CheckAndInferReturnAny => {
-                    let implicit_return = self.implicit_return(&body, func_name);
+                    let implicit_return = self.implicit_return(&body, func_name, parameters);
                     let (yields_and_returns, self_assignments, unused_parameters, unused_variables) =
                         self.function_body_scope(
                             parameters,
@@ -657,7 +738,7 @@ impl<'a> BindingsBuilder<'a> {
                     self_assignments
                 }
                 UntypedDefBehavior::CheckAndInferReturnType => {
-                    let implicit_return = self.implicit_return(&body, func_name);
+                    let implicit_return = self.implicit_return(&body, func_name, parameters);
                     let (yields_and_returns, self_assignments, unused_parameters, unused_variables) =
                         self.function_body_scope(
                             parameters,
@@ -770,6 +851,47 @@ impl<'a> BindingsBuilder<'a> {
     }
 }
 
+fn extract_isinstance_test(test: &Expr) -> Option<(Name, SmallSet<Name>)> {
+    let Expr::Call(ExprCall {
+        func, arguments, ..
+    }) = test
+    else {
+        return None;
+    };
+    let Expr::Name(func_name) = func.as_ref() else {
+        return None;
+    };
+    if func_name.id.as_str() != "isinstance" {
+        return None;
+    }
+    if !arguments.keywords.is_empty() || arguments.args.len() != 2 {
+        return None;
+    }
+    let Expr::Name(subject) = &arguments.args[0] else {
+        return None;
+    };
+    let mut types = SmallSet::new();
+    match &arguments.args[1] {
+        Expr::Name(name) => {
+            types.insert(name.id.clone());
+        }
+        Expr::Tuple(tuple) => {
+            for elt in &tuple.elts {
+                let Expr::Name(name) = elt else {
+                    return None;
+                };
+                types.insert(name.id.clone());
+            }
+        }
+        _ => return None,
+    }
+    if types.is_empty() {
+        None
+    } else {
+        Some((subject.id.clone(), types))
+    }
+}
+
 /// Given the body of a function, what are the potential expressions that
 /// could be the last ones to be executed, where the function then falls off the end.
 ///
@@ -779,8 +901,14 @@ impl<'a> BindingsBuilder<'a> {
 fn function_last_expressions<'a>(
     x: &'a [Stmt],
     sys_info: &SysInfo,
+    constrained_typevars: Option<&ConstrainedTypeVarParams>,
 ) -> Option<Vec<(LastStmt, &'a Expr)>> {
-    fn f<'a>(sys_info: &SysInfo, x: &'a [Stmt], res: &mut Vec<(LastStmt, &'a Expr)>) -> Option<()> {
+    fn f<'a>(
+        sys_info: &SysInfo,
+        x: &'a [Stmt],
+        res: &mut Vec<(LastStmt, &'a Expr)>,
+        constrained_typevars: Option<&ConstrainedTypeVarParams>,
+    ) -> Option<()> {
         fn loop_body_has_break_statement(statement: &Stmt, has_break: &mut bool) {
             match statement {
                 Stmt::Break(_) => {
@@ -801,7 +929,7 @@ fn function_last_expressions<'a>(
                 for y in &x.items {
                     res.push((LastStmt::With(kind), &y.context_expr));
                 }
-                f(sys_info, &x.body, res)?;
+                f(sys_info, &x.body, res, constrained_typevars)?;
             }
             Stmt::While(x) => {
                 // Infinite loops with no breaks cannot fall through
@@ -822,17 +950,50 @@ fn function_last_expressions<'a>(
                 if has_break || x.orelse.is_empty() {
                     return None;
                 }
-                f(sys_info, &x.orelse, res)?;
+                f(sys_info, &x.orelse, res, constrained_typevars)?;
             }
             Stmt::If(x) => {
                 let mut last_test = None;
+                let mut chain_var = None;
+                let mut covered = SmallSet::new();
+                let mut chain_valid = constrained_typevars.is_some();
                 for (test, body) in sys_info.pruned_if_branches(x) {
                     last_test = test;
-                    f(sys_info, body, res)?;
+                    if let (Some(test), true) = (test, chain_valid) {
+                        // Special-case isinstance chains over constrained TypeVars to avoid
+                        // false missing-return errors when all constraints are covered.
+                        if let Some((var, types)) = extract_isinstance_test(test) {
+                            if let Some(existing) = &chain_var {
+                                if existing != &var {
+                                    chain_valid = false;
+                                }
+                            } else {
+                                chain_var = Some(var);
+                            }
+                            if chain_valid {
+                                for ty in types {
+                                    covered.insert(ty);
+                                }
+                            }
+                        } else {
+                            chain_valid = false;
+                        }
+                    }
+                    f(sys_info, body, res, constrained_typevars)?;
                 }
                 if last_test.is_some() {
-                    // The final `if` can fall through, so the `if` itself might be the last statement.
-                    return None;
+                    let mut exhaustive = false;
+                    if chain_valid
+                        && let (Some(var), Some(constrained_typevars)) =
+                            (&chain_var, constrained_typevars)
+                        && let Some(constraints) = constrained_typevars.get(var)
+                    {
+                        exhaustive = constraints.iter().all(|c| covered.contains(c));
+                    }
+                    if !exhaustive {
+                        // The final `if` can fall through, so the `if` itself might be the last statement.
+                        return None;
+                    }
                 }
             }
             Stmt::Try(x) => {
@@ -843,16 +1004,18 @@ fn function_last_expressions<'a>(
                         .iter()
                         .any(|stmt| matches!(stmt, Stmt::Return(_)))
                 {
-                    f(sys_info, &x.finalbody, res)?;
+                    f(sys_info, &x.finalbody, res, constrained_typevars)?;
                 } else {
                     if x.orelse.is_empty() {
-                        f(sys_info, &x.body, res)?;
+                        f(sys_info, &x.body, res, constrained_typevars)?;
                     } else {
-                        f(sys_info, &x.orelse, res)?;
+                        f(sys_info, &x.orelse, res, constrained_typevars)?;
                     }
                     for handler in &x.handlers {
                         match handler {
-                            ExceptHandler::ExceptHandler(x) => f(sys_info, &x.body, res)?,
+                            ExceptHandler::ExceptHandler(x) => {
+                                f(sys_info, &x.body, res, constrained_typevars)?
+                            }
                         }
                     }
                     // If we don't have a matching handler, we raise an exception, which is fine.
@@ -861,7 +1024,7 @@ fn function_last_expressions<'a>(
             Stmt::Match(x) => {
                 let mut exhaustive = false;
                 for case in x.cases.iter() {
-                    f(sys_info, &case.body, res)?;
+                    f(sys_info, &case.body, res, constrained_typevars)?;
                     if case.pattern.is_wildcard() || case.pattern.is_irrefutable() {
                         exhaustive = true;
                         break;
@@ -877,7 +1040,7 @@ fn function_last_expressions<'a>(
     }
 
     let mut res = Vec::new();
-    f(sys_info, x, &mut res)?;
+    f(sys_info, x, &mut res, constrained_typevars)?;
     Some(res)
 }
 

pyrefly/lib/solver/subset.rs

@@ -991,7 +991,7 @@ impl<'a, Ans: LookupAnswer> Subset<'a, Ans> {
             }
             (t1, Type::Quantified(q)) => match q.restriction() {
                 // This only works for constraints and not bounds, because a TypeVar must resolve to exactly one of its constraints.
-                Restriction::Constraints(constraints) => all(constraints.iter(), |constraint| {
+                Restriction::Constraints(constraints) => any(constraints.iter(), |constraint| {
                     self.is_subset_eq(t1, constraint)
                 }),
                 _ => Err(SubsetError::Other),

pyrefly/lib/test/generic_restrictions.rs

@@ -523,7 +523,7 @@ testcase!(
     test_add_with_constraints,
     r#"
 def add[T: (int, str)](x: T, y: T) -> T:
-    return x + y # E: `+` is not supported # E: `+` is not supported # E: `int | Unknown` is not assignable to declared return type `T`
+    return x + y # E: `+` is not supported # E: `+` is not supported
     "#,
 );
 

pyrefly/lib/test/returns.rs

@@ -151,6 +151,21 @@ def f(b: bool) -> int:  # E: Function declared to return `int`, but one or more
 "#,
 );
 
+testcase!(
+    test_return_constrained_typevar_isinstance,
+    r#"
+from typing import TypeVar
+
+T = TypeVar("T", int, str)
+
+def f(x: T) -> T:
+    if isinstance(x, int):
+        return x
+    elif isinstance(x, str):
+        return x
+"#,
+);
+
 testcase!(
     test_return_if_no_else_none,
     r#"