From 24cd592a1d6d5ad8ec4e0cb0e836f94771d12afd Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Thu, 14 Nov 2024 21:24:12 +0100
Subject: [PATCH] Avoid module lookup for known classes when possible (#14343)

---
 crates/red_knot_python_semantic/src/types.rs  | 19 ++++++++++---------
 .../src/types/infer.rs                        |  4 +---
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 2b23469408d2f..4e26fea03a64e 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -6,6 +6,13 @@ use itertools::Itertools;
 use ruff_db::files::File;
 use ruff_python_ast as ast;
 
+pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
+pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
+pub(crate) use self::display::TypeArrayDisplay;
+pub(crate) use self::infer::{
+    infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types,
+};
+use crate::module_resolver::file_to_module;
 use crate::semantic_index::ast_ids::HasScopedAstId;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId};
@@ -22,13 +29,6 @@ use crate::types::mro::{ClassBase, Mro, MroError, MroIterator};
 use crate::types::narrow::narrowing_constraint;
 use crate::{Db, FxOrderSet, Module, Program};
 
-pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
-pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
-pub(crate) use self::display::TypeArrayDisplay;
-pub(crate) use self::infer::{
-    infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types,
-};
-
 mod builder;
 mod diagnostic;
 mod display;
@@ -1720,7 +1720,7 @@ impl<'db> KnownClass {
         }
     }
 
-    pub fn try_from_module(module: &Module, class_name: &str) -> Option<Self> {
+    pub fn try_from_file(db: &dyn Db, file: File, class_name: &str) -> Option<Self> {
         // Note: if this becomes hard to maintain (as rust can't ensure at compile time that all
         // variants of `Self` are covered), we might use a macro (in-house or dependency)
         // See: https://stackoverflow.com/q/39070244
@@ -1747,7 +1747,8 @@ impl<'db> KnownClass {
             _ => return None,
         };
 
-        candidate.check_module(module).then_some(candidate)
+        let module = file_to_module(db, file)?;
+        candidate.check_module(&module).then_some(candidate)
     }
 
     /// Return `true` if the module of `self` matches `module_name`
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 443835d90de04..578f56fbcf4c0 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -1041,9 +1041,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             .node_scope(NodeWithScopeRef::Class(class_node))
             .to_scope_id(self.db, self.file);
 
-        let maybe_known_class = file_to_module(self.db, self.file)
-            .as_ref()
-            .and_then(|module| KnownClass::try_from_module(module, name.as_str()));
+        let maybe_known_class = KnownClass::try_from_file(self.db, self.file, name);
 
         let class = Class::new(self.db, &*name.id, body_scope, maybe_known_class);
         let class_ty = Type::class_literal(class);