[red-knot] Infer Literal types from comparisons with sys.version_info (#14244)

crates/red_knot_python_semantic/src/stdlib.rs

@@ -14,6 +14,7 @@ pub(crate) enum CoreStdlibModule {
     Typeshed,
     TypingExtensions,
     Typing,
+    Sys,
 }
 
 impl CoreStdlibModule {
@@ -24,6 +25,7 @@ impl CoreStdlibModule {
             Self::Typing => "typing",
             Self::Typeshed => "_typeshed",
             Self::TypingExtensions => "typing_extensions",
+            Self::Sys => "sys",
         }
     }
 

crates/red_knot_python_semantic/src/types.rs

@@ -20,7 +20,7 @@ use crate::symbol::{Boundness, Symbol};
 use crate::types::diagnostic::TypeCheckDiagnosticsBuilder;
 use crate::types::mro::{ClassBase, Mro, MroError, MroIterator};
 use crate::types::narrow::narrowing_constraint;
-use crate::{Db, FxOrderSet, HasTy, Module, SemanticModel};
+use crate::{Db, FxOrderSet, HasTy, Module, Program, SemanticModel};
 
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
@@ -991,7 +991,9 @@ impl<'db> Type<'db> {
                 .all(|elem| elem.is_single_valued(db)),
 
             Type::Instance(InstanceType { class }) => match class.known(db) {
-                Some(KnownClass::NoneType | KnownClass::NoDefaultType) => true,
+                Some(
+                    KnownClass::NoneType | KnownClass::NoDefaultType | KnownClass::VersionInfo,
+                ) => true,
                 Some(
                     KnownClass::Bool
                     | KnownClass::Object
@@ -1081,9 +1083,18 @@ impl<'db> Type<'db> {
             Type::ClassLiteral(class_ty) => class_ty.member(db, name),
             Type::SubclassOf(subclass_of_ty) => subclass_of_ty.member(db, name),
             Type::KnownInstance(known_instance) => known_instance.member(db, name),
-            Type::Instance(_) => {
-                // TODO MRO? get_own_instance_member, get_instance_member
-                Type::Todo.into()
+            Type::Instance(InstanceType { class }) => {
+                let ty = match (class.known(db), name) {
+                    (Some(KnownClass::VersionInfo), "major") => {
+                        Type::IntLiteral(Program::get(db).target_version(db).major.into())
+                    }
+                    (Some(KnownClass::VersionInfo), "minor") => {
+                        Type::IntLiteral(Program::get(db).target_version(db).minor.into())
+                    }
+                    // TODO MRO? get_own_instance_member, get_instance_member
+                    _ => Type::Todo,
+                };
+                ty.into()
             }
             Type::Union(union) => {
                 let mut builder = UnionBuilder::new(db);
@@ -1418,6 +1429,39 @@ impl<'db> Type<'db> {
         KnownClass::NoneType.to_instance(db)
     }
 
+    /// Return the type of `tuple(sys.version_info)`.
+    ///
+    /// This is not exactly the type that `sys.version_info` has at runtime,
+    /// but it's a useful fallback for us in order to infer `Literal` types from `sys.version_info` comparisons.
+    fn version_info_tuple(db: &'db dyn Db) -> Self {
+        let target_version = Program::get(db).target_version(db);
+        let int_instance_ty = KnownClass::Int.to_instance(db);
+
+        // TODO: just grab this type from typeshed (it's a `sys._ReleaseLevel` type alias there)
+        let release_level_ty = {
+            let elements: Box<[Type<'db>]> = ["alpha", "beta", "candidate", "final"]
+                .iter()
+                .map(|level| Type::string_literal(db, level))
+                .collect();
+
+            // For most unions, it's better to go via `UnionType::from_elements` or use `UnionBuilder`;
+            // those techniques ensure that union elements are deduplicated and unions are eagerly simplified
+            // into other types where necessary. Here, however, we know that there are no duplicates
+            // in this union, so it's probably more efficient to use `UnionType::new()` directly.
+            Type::Union(UnionType::new(db, elements))
+        };
+
+        let version_info_elements = &[
+            Type::IntLiteral(target_version.major.into()),
+            Type::IntLiteral(target_version.minor.into()),
+            int_instance_ty,
+            release_level_ty,
+            int_instance_ty,
+        ];
+
+        Self::tuple(db, version_info_elements)
+    }
+
     /// Given a type that is assumed to represent an instance of a class,
     /// return a type that represents that class itself.
     #[must_use]
@@ -1541,6 +1585,8 @@ pub enum KnownClass {
     SpecialForm,
     TypeVar,
     NoDefaultType,
+    // sys
+    VersionInfo,
 }
 
 impl<'db> KnownClass {
@@ -1565,6 +1611,12 @@ impl<'db> KnownClass {
             Self::SpecialForm => "_SpecialForm",
             Self::TypeVar => "TypeVar",
             Self::NoDefaultType => "_NoDefaultType",
+            // This is the name the type of `sys.version_info` has in typeshed,
+            // which is different to what `type(sys.version_info).__name__` is at runtime.
+            // (At runtime, `type(sys.version_info).__name__ == "version_info"`,
+            // which is impossible to replicate in the stubs since the sole instance of the class
+            // also has that name in the `sys` module.)
+            Self::VersionInfo => "_version_info",
         }
     }
 
@@ -1591,6 +1643,7 @@ impl<'db> KnownClass {
             | Self::Set
             | Self::Dict
             | Self::Slice => CoreStdlibModule::Builtins,
+            Self::VersionInfo => CoreStdlibModule::Sys,
             Self::GenericAlias | Self::ModuleType | Self::FunctionType => CoreStdlibModule::Types,
             Self::NoneType => CoreStdlibModule::Typeshed,
             Self::SpecialForm | Self::TypeVar => CoreStdlibModule::Typing,
@@ -1607,7 +1660,7 @@ impl<'db> KnownClass {
     const fn is_singleton(self) -> bool {
         // TODO there are other singleton types (EllipsisType, NotImplementedType)
         match self {
-            Self::NoneType | Self::NoDefaultType => true,
+            Self::NoneType | Self::NoDefaultType | Self::VersionInfo => true,
             Self::Bool
             | Self::Object
             | Self::Bytes
@@ -1651,13 +1704,14 @@ impl<'db> KnownClass {
             "FunctionType" => Self::FunctionType,
             "_SpecialForm" => Self::SpecialForm,
             "_NoDefaultType" => Self::NoDefaultType,
+            "_version_info" => Self::VersionInfo,
             _ => return None,
         };
 
         candidate.check_module(module).then_some(candidate)
     }
 
-    /// Private method checking if known class can be defined in the given module.
+    /// Return `true` if the module of `self` matches `module_name`
     fn check_module(self, module: &Module) -> bool {
         if !module.search_path().is_standard_library() {
             return false;
@@ -1677,6 +1731,7 @@ impl<'db> KnownClass {
             | Self::Slice
             | Self::GenericAlias
             | Self::ModuleType
+            | Self::VersionInfo
             | Self::FunctionType => module.name() == self.canonical_module().as_str(),
             Self::NoneType => matches!(module.name().as_str(), "_typeshed" | "types"),
             Self::SpecialForm | Self::TypeVar | Self::NoDefaultType => {

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1628,8 +1628,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         let mut annotation_ty = self.infer_annotation_expression(annotation);
 
-        // If the declared variable is annotated with _SpecialForm class then we treat it differently
-        // by assigning the known field to the instance.
+        // Handle various singletons.
         if let Type::Instance(InstanceType { class }) = annotation_ty {
             if class.is_known(self.db, KnownClass::SpecialForm) {
                 if let Some(name_expr) = target.as_name_expr() {
@@ -3531,6 +3530,16 @@ impl<'db> TypeInferenceBuilder<'db> {
             (_, Type::BytesLiteral(_)) => {
                 self.infer_binary_type_comparison(left, op, KnownClass::Bytes.to_instance(self.db))
             }
+            (Type::Tuple(_), Type::Instance(InstanceType { class }))
+                if class.is_known(self.db, KnownClass::VersionInfo) =>
+            {
+                self.infer_binary_type_comparison(left, op, Type::version_info_tuple(self.db))
+            }
+            (Type::Instance(InstanceType { class }), Type::Tuple(_))
+                if class.is_known(self.db, KnownClass::VersionInfo) =>
+            {
+                self.infer_binary_type_comparison(Type::version_info_tuple(self.db), op, right)
+            }
             (Type::Tuple(lhs), Type::Tuple(rhs)) => {
                 // Note: This only works on heterogeneous tuple types.
                 let lhs_elements = lhs.elements(self.db);
@@ -3713,6 +3722,16 @@ impl<'db> TypeInferenceBuilder<'db> {
         slice_ty: Type<'db>,
     ) -> Type<'db> {
         match (value_ty, slice_ty) {
+            (
+                Type::Instance(InstanceType { class }),
+                Type::IntLiteral(_) | Type::BooleanLiteral(_) | Type::SliceLiteral(_),
+            ) if class.is_known(self.db, KnownClass::VersionInfo) => self
+                .infer_subscript_expression_types(
+                    value_node,
+                    Type::version_info_tuple(self.db),
+                    slice_ty,
+                ),
+
             // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"`
             (Type::Tuple(tuple_ty), Type::IntLiteral(int)) if i32::try_from(int).is_ok() => {
                 let elements = tuple_ty.elements(self.db);