Simplify argument handling

astral-sh · Jan 7, 2025 · ccb639c · ccb639c
1 parent 74e8117
commit ccb639c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 14 deletions.
diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
@@ -3249,6 +3249,27 @@ impl KnownFunction {
             _ => None,
         }
     }
+
+    /// Whether or not a particular function takes type expression as arguments, i.e.
+    /// should the argument of a call like `f(int)` be interpreted as the type int (true)
+    /// or as the type of the expression `int`, i.e. Literal[int] (false).
+    fn takes_type_expression_arguments(self) -> bool {
+        match self {
+            KnownFunction::IsEquivalentTo
+            | KnownFunction::IsSubtypeOf
+            | KnownFunction::IsAssignableTo
+            | KnownFunction::IsDisjointFrom
+            | KnownFunction::IsFullyStatic
+            | KnownFunction::IsSingleton
+            | KnownFunction::IsSingleValued => true,
+            KnownFunction::ConstraintFunction(_)
+            | KnownFunction::RevealType
+            | KnownFunction::Len
+            | KnownFunction::Final
+            | KnownFunction::NoTypeCheck
+            | KnownFunction::StaticAssert => false,
+        }
+    }
 }
 
 #[salsa::interned]

diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -3049,20 +3049,11 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         let function_type = self.infer_expression(func);
 
-        let infer_arguments_as_type_expressions = matches!(
-            function_type
-                .into_function_literal()
-                .and_then(|f| f.known(self.db())),
-            Some(
-                KnownFunction::IsEquivalentTo
-                    | KnownFunction::IsSubtypeOf
-                    | KnownFunction::IsAssignableTo
-                    | KnownFunction::IsDisjointFrom
-                    | KnownFunction::IsFullyStatic
-                    | KnownFunction::IsSingleton
-                    | KnownFunction::IsSingleValued
-            ),
-        );
+        let infer_arguments_as_type_expressions = function_type
+            .into_function_literal()
+            .and_then(|f| f.known(self.db()))
+            .is_some_and(|f| f.takes_type_expression_arguments());
+
         let call_arguments = self.infer_arguments(arguments, infer_arguments_as_type_expressions);
         function_type
             .call(self.db(), &call_arguments)