[ty] Add special-cased inference for __import__(name) and importlib.import_module(name) (#19008)

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -154,6 +154,8 @@ pub enum KnownModule {
     #[strum(serialize = "_typeshed._type_checker_internals")]
     TypeCheckerInternals,
     TyExtensions,
+    #[strum(serialize = "importlib")]
+    ImportLib,
 }
 
 impl KnownModule {
@@ -172,6 +174,7 @@ impl KnownModule {
             Self::Inspect => "inspect",
             Self::TypeCheckerInternals => "_typeshed._type_checker_internals",
             Self::TyExtensions => "ty_extensions",
+            Self::ImportLib => "importlib",
         }
     }
 
@@ -210,6 +213,10 @@ impl KnownModule {
     pub const fn is_enum(self) -> bool {
         matches!(self, Self::Enum)
     }
+
+    pub const fn is_importlib(self) -> bool {
+        matches!(self, Self::ImportLib)
+    }
 }
 
 impl std::fmt::Display for KnownModule {

crates/ty_python_semantic/src/types/call/bind.rs

@@ -740,7 +740,7 @@ impl<'db> Bindings<'db> {
 
                         Some(KnownFunction::Override) => {
                             // TODO: This can be removed once we understand legacy generics because the
-                            // typeshed definition for `typing.overload` is an identity function.
+                            // typeshed definition for `typing.override` is an identity function.
                             if let [Some(ty)] = overload.parameter_types() {
                                 overload.set_return_type(*ty);
                             }
@@ -756,7 +756,7 @@ impl<'db> Bindings<'db> {
 
                         Some(KnownFunction::Final) => {
                             // TODO: This can be removed once we understand legacy generics because the
-                            // typeshed definition for `abc.abstractmethod` is an identity function.
+                            // typeshed definition for `typing.final` is an identity function.
                             if let [Some(ty)] = overload.parameter_types() {
                                 overload.set_return_type(*ty);
                             }

crates/ty_python_semantic/src/types/function.rs

@@ -77,7 +77,7 @@ use crate::types::{
     BoundMethodType, CallableType, DynamicType, KnownClass, Type, TypeMapping, TypeRelation,
     TypeVarInstance,
 };
-use crate::{Db, FxOrderSet};
+use crate::{Db, FxOrderSet, ModuleName, resolve_module};
 
 /// A collection of useful spans for annotating functions.
 ///
@@ -867,6 +867,12 @@ pub enum KnownFunction {
     Len,
     /// `builtins.repr`
     Repr,
+    /// `builtins.__import__`, which returns the top-level module.
+    #[strum(serialize = "__import__")]
+    DunderImport,
+    /// `importlib.import_module`, which returns the submodule.
+    ImportModule,
+
     /// `typing(_extensions).final`
     Final,
 
@@ -951,9 +957,12 @@ impl KnownFunction {
     /// Return `true` if `self` is defined in `module` at runtime.
     const fn check_module(self, module: KnownModule) -> bool {
         match self {
-            Self::IsInstance | Self::IsSubclass | Self::HasAttr | Self::Len | Self::Repr => {
-                module.is_builtins()
-            }
+            Self::IsInstance
+            | Self::IsSubclass
+            | Self::HasAttr
+            | Self::Len
+            | Self::Repr
+            | Self::DunderImport => module.is_builtins(),
             Self::AssertType
             | Self::AssertNever
             | Self::Cast
@@ -987,48 +996,45 @@ impl KnownFunction {
             | Self::DunderAllNames
             | Self::StaticAssert
             | Self::AllMembers => module.is_ty_extensions(),
+            Self::ImportModule => module.is_importlib(),
         }
     }
 
     /// Evaluate a call to this known function, and emit any diagnostics that are necessary
     /// as a result of the call.
-    pub(super) fn check_call(
+    pub(super) fn check_call<'db>(
         self,
-        context: &InferContext,
-        parameter_types: &[Option<Type<'_>>],
+        context: &InferContext<'db, '_>,
+        parameter_types: &[Option<Type<'db>>],
         call_expression: &ast::ExprCall,
-    ) {
+        file: File,
+    ) -> Option<Type<'db>> {
         let db = context.db();
 
         match self {
             KnownFunction::RevealType => {
                 let [Some(revealed_type)] = parameter_types else {
-                    return;
-                };
-                let Some(builder) =
-                    context.report_diagnostic(DiagnosticId::RevealedType, Severity::Info)
-                else {
-                    return;
+                    return None;
                 };
+                let builder =
+                    context.report_diagnostic(DiagnosticId::RevealedType, Severity::Info)?;
                 let mut diag = builder.into_diagnostic("Revealed type");
                 let span = context.span(&call_expression.arguments.args[0]);
                 diag.annotate(
                     Annotation::primary(span)
                         .message(format_args!("`{}`", revealed_type.display(db))),
                 );
+                None
             }
             KnownFunction::AssertType => {
                 let [Some(actual_ty), Some(asserted_ty)] = parameter_types else {
-                    return;
+                    return None;
                 };
 
                 if actual_ty.is_equivalent_to(db, *asserted_ty) {
-                    return;
+                    return None;
                 }
-                let Some(builder) = context.report_lint(&TYPE_ASSERTION_FAILURE, call_expression)
-                else {
-                    return;
-                };
+                let builder = context.report_lint(&TYPE_ASSERTION_FAILURE, call_expression)?;
 
                 let mut diagnostic = builder.into_diagnostic(format_args!(
                     "Argument does not have asserted type `{}`",
@@ -1048,18 +1054,17 @@ impl KnownFunction {
                     asserted_type = asserted_ty.display(db),
                     inferred_type = actual_ty.display(db),
                 ));
+
+                None
             }
             KnownFunction::AssertNever => {
                 let [Some(actual_ty)] = parameter_types else {
-                    return;
+                    return None;
                 };
                 if actual_ty.is_equivalent_to(db, Type::Never) {
-                    return;
+                    return None;
                 }
-                let Some(builder) = context.report_lint(&TYPE_ASSERTION_FAILURE, call_expression)
-                else {
-                    return;
-                };
+                let builder = context.report_lint(&TYPE_ASSERTION_FAILURE, call_expression)?;
 
                 let mut diagnostic =
                     builder.into_diagnostic("Argument does not have asserted type `Never`");
@@ -1074,10 +1079,12 @@ impl KnownFunction {
                     "`Never` and `{inferred_type}` are not equivalent types",
                     inferred_type = actual_ty.display(db),
                 ));
+
+                None
             }
             KnownFunction::StaticAssert => {
                 let [Some(parameter_ty), message] = parameter_types else {
-                    return;
+                    return None;
                 };
                 let truthiness = match parameter_ty.try_bool(db) {
                     Ok(truthiness) => truthiness,
@@ -1097,16 +1104,13 @@ impl KnownFunction {
 
                         err.report_diagnostic(context, condition);
 
-                        return;
+                        return None;
                     }
                 };
 
-                let Some(builder) = context.report_lint(&STATIC_ASSERT_ERROR, call_expression)
-                else {
-                    return;
-                };
+                let builder = context.report_lint(&STATIC_ASSERT_ERROR, call_expression)?;
                 if truthiness.is_always_true() {
-                    return;
+                    return None;
                 }
                 if let Some(message) = message
                     .and_then(Type::into_string_literal)
@@ -1129,10 +1133,12 @@ impl KnownFunction {
                         parameter_ty = parameter_ty.display(db)
                     ));
                 }
+
+                None
             }
             KnownFunction::Cast => {
                 let [Some(casted_type), Some(source_type)] = parameter_types else {
-                    return;
+                    return None;
                 };
                 let contains_unknown_or_todo =
                     |ty| matches!(ty, Type::Dynamic(dynamic) if dynamic != DynamicType::Any);
@@ -1140,34 +1146,31 @@ impl KnownFunction {
                     && !casted_type.any_over_type(db, &|ty| contains_unknown_or_todo(ty))
                     && !source_type.any_over_type(db, &|ty| contains_unknown_or_todo(ty))
                 {
-                    let Some(builder) = context.report_lint(&REDUNDANT_CAST, call_expression)
-                    else {
-                        return;
-                    };
+                    let builder = context.report_lint(&REDUNDANT_CAST, call_expression)?;
                     builder.into_diagnostic(format_args!(
                         "Value is already of type `{}`",
                         casted_type.display(db),
                     ));
                 }
+                None
             }
             KnownFunction::GetProtocolMembers => {
                 let [Some(Type::ClassLiteral(class))] = parameter_types else {
-                    return;
+                    return None;
                 };
                 if class.is_protocol(db) {
-                    return;
+                    return None;
                 }
                 report_bad_argument_to_get_protocol_members(context, call_expression, *class);
+                None
             }
             KnownFunction::IsInstance | KnownFunction::IsSubclass => {
                 let [_, Some(Type::ClassLiteral(class))] = parameter_types else {
-                    return;
-                };
-                let Some(protocol_class) = class.into_protocol_class(db) else {
-                    return;
+                    return None;
                 };
+                let protocol_class = class.into_protocol_class(db)?;
                 if protocol_class.is_runtime_checkable(db) {
-                    return;
+                    return None;
                 }
                 report_runtime_check_against_non_runtime_checkable_protocol(
                     context,
@@ -1175,8 +1178,35 @@ impl KnownFunction {
                     protocol_class,
                     self,
                 );
+                None
             }
-            _ => {}
+            known @ (KnownFunction::DunderImport | KnownFunction::ImportModule) => {
+                let [Some(Type::StringLiteral(full_module_name)), rest @ ..] = parameter_types
+                else {
+                    return None;
+                };
+
+                if rest.iter().any(Option::is_some) {
+                    return None;
+                }
+
+                let module_name = full_module_name.value(db);
+
+                if known == KnownFunction::DunderImport && module_name.contains('.') {
+                    // `__import__("collections.abc")` returns the `collections` module.
+                    // `importlib.import_module("collections.abc")` returns the `collections.abc` module.
+                    // ty doesn't have a way to represent the return type of the former yet.
+                    // https://github.com/astral-sh/ruff/pull/19008#discussion_r2173481311
+                    return None;
+                }
+
+                let module_name = ModuleName::new(module_name)?;
+                let module = resolve_module(db, &module_name)?;
+
+                Some(Type::module_literal(db, file, &module))
+            }
+
+            _ => None,
         }
     }
 }
@@ -1201,7 +1231,8 @@ pub(crate) mod tests {
                 | KnownFunction::Repr
                 | KnownFunction::IsInstance
                 | KnownFunction::HasAttr
-                | KnownFunction::IsSubclass => KnownModule::Builtins,
+                | KnownFunction::IsSubclass
+                | KnownFunction::DunderImport => KnownModule::Builtins,
 
                 KnownFunction::AbstractMethod => KnownModule::Abc,
 
@@ -1234,6 +1265,8 @@ pub(crate) mod tests {
                 | KnownFunction::TopMaterialization
                 | KnownFunction::BottomMaterialization
                 | KnownFunction::AllMembers => KnownModule::TyExtensions,
+
+                KnownFunction::ImportModule => KnownModule::ImportLib,
             };
 
             let function_definition = known_module_symbol(&db, module, function_name)

crates/ty_python_semantic/src/types/infer.rs

@@ -5388,11 +5388,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 match binding_type {
                     Type::FunctionLiteral(function_literal) => {
                         if let Some(known_function) = function_literal.known(self.db()) {
-                            known_function.check_call(
+                            if let Some(return_type) = known_function.check_call(
                                 &self.context,
                                 overload.parameter_types(),
                                 call_expression,
-                            );
+                                self.file(),
+                            ) {
+                                overload.set_return_type(return_type);
+                            }
                         }
                     }