`UnionType::known`

This commit is contained in:
Jack O'Connor 2025-12-11 13:46:06 -08:00
parent 2b109dfeaa
commit d90d33001e
3 changed files with 53 additions and 64 deletions

View File

@ -13916,40 +13916,13 @@ impl<'db> UnionType<'db> {
ConstraintSet::from(sorted_self == other.normalized(db)) ConstraintSet::from(sorted_self == other.normalized(db))
} }
/// Returns true if this union is equivalent to `int | float`, which is what `float` expands
/// into in type position.
pub(crate) fn is_int_float(self, db: &'db dyn Db) -> bool {
let elements = self.elements(db);
if elements.len() != 2 {
return false;
}
let mut has_int = false;
let mut has_float = false;
for element in elements {
if let Type::NominalInstance(nominal) = element
&& let Some(known) = nominal.known_class(db)
{
match known {
KnownClass::Int => has_int = true,
KnownClass::Float => has_float = true,
_ => {}
}
}
}
has_int && has_float
}
/// Returns true if this union is equivalent to `int | float | complex`, which is what /// Returns true if this union is equivalent to `int | float | complex`, which is what
/// `complex` expands into in type position. /// `complex` expands into in type position.
pub(crate) fn is_int_float_complex(self, db: &'db dyn Db) -> bool { pub(crate) fn known(self, db: &'db dyn Db) -> Option<KnownUnion> {
let elements = self.elements(db);
if elements.len() != 3 {
return false;
}
let mut has_int = false; let mut has_int = false;
let mut has_float = false; let mut has_float = false;
let mut has_complex = false; let mut has_complex = false;
for element in elements { for element in self.elements(db) {
if let Type::NominalInstance(nominal) = element if let Type::NominalInstance(nominal) = element
&& let Some(known) = nominal.known_class(db) && let Some(known) = nominal.known_class(db)
{ {
@ -13957,12 +13930,24 @@ impl<'db> UnionType<'db> {
KnownClass::Int => has_int = true, KnownClass::Int => has_int = true,
KnownClass::Float => has_float = true, KnownClass::Float => has_float = true,
KnownClass::Complex => has_complex = true, KnownClass::Complex => has_complex = true,
_ => {} _ => return None,
}
} else {
return None;
}
}
match (has_int, has_float, has_complex) {
(true, true, false) => Some(KnownUnion::Float),
(true, true, true) => Some(KnownUnion::Complex),
_ => None,
} }
} }
} }
has_int && has_float && has_complex
} #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum KnownUnion {
Float, // `int | float`
Complex, // `int | float | complex`
} }
#[salsa::interned(debug, heap_size=IntersectionType::heap_size)] #[salsa::interned(debug, heap_size=IntersectionType::heap_size)]

View File

@ -104,13 +104,13 @@ use crate::types::visitor::any_over_type;
use crate::types::{ use crate::types::{
BoundTypeVarInstance, CallDunderError, CallableBinding, CallableType, CallableTypeKind, BoundTypeVarInstance, CallDunderError, CallableBinding, CallableType, CallableTypeKind,
ClassLiteral, ClassType, DataclassParams, DynamicType, InternedType, IntersectionBuilder, ClassLiteral, ClassType, DataclassParams, DynamicType, InternedType, IntersectionBuilder,
IntersectionType, KnownClass, KnownInstanceType, LintDiagnosticGuard, MemberLookupPolicy, IntersectionType, KnownClass, KnownInstanceType, KnownUnion, LintDiagnosticGuard,
MetaclassCandidate, PEP695TypeAliasType, ParamSpecAttrKind, Parameter, ParameterForm, MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, ParamSpecAttrKind, Parameter,
Parameters, Signature, SpecialFormType, SubclassOfType, TrackedConstraintSet, Truthiness, Type, ParameterForm, Parameters, Signature, SpecialFormType, SubclassOfType, TrackedConstraintSet,
TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers, TypeVarBoundOrConstraints, Truthiness, Type, TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers,
TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarIdentity, TypeVarBoundOrConstraints, TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation,
TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder, UnionType, TypeVarIdentity, TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder,
UnionTypeInstance, binding_type, infer_scope_types, todo_type, UnionType, UnionTypeInstance, binding_type, infer_scope_types, todo_type,
}; };
use crate::types::{CallableTypes, overrides}; use crate::types::{CallableTypes, overrides};
use crate::types::{ClassBase, add_inferred_python_version_hint_to_diagnostic}; use crate::types::{ClassBase, add_inferred_python_version_hint_to_diagnostic};
@ -5629,35 +5629,34 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
// Infer the deferred base type of a NewType. // Infer the deferred base type of a NewType.
fn infer_newtype_assignment_deferred(&mut self, arguments: &ast::Arguments) { fn infer_newtype_assignment_deferred(&mut self, arguments: &ast::Arguments) {
match self.infer_type_expression(&arguments.args[1]) { let inferred = self.infer_type_expression(&arguments.args[1]);
Type::NominalInstance(_) | Type::NewTypeInstance(_) => {} match inferred {
Type::NominalInstance(_) | Type::NewTypeInstance(_) => return,
// There are exactly two union types allowed as bases for NewType: `int | float` and // There are exactly two union types allowed as bases for NewType: `int | float` and
// `int | float | complex`. These are allowed because that's what `float` and `complex` // `int | float | complex`. These are allowed because that's what `float` and `complex`
// expand into in type position. We don't currently ask whether the union was implicit // expand into in type position. We don't currently ask whether the union was implicit
// or explicit, so the explicit version is also allowed. // or explicit, so the explicit version is also allowed.
Type::Union(union_ty) Type::Union(union_ty) => match union_ty.known(self.db()) {
if union_ty.is_int_float(self.db()) || union_ty.is_int_float_complex(self.db()) => { Some(KnownUnion::Float) | Some(KnownUnion::Complex) => return,
} _ => {}
},
// `Unknown` is likely to be the result of an unresolved import or a typo, which will // `Unknown` is likely to be the result of an unresolved import or a typo, which will
// already get a diagnostic, so don't pile on an extra diagnostic here. // already get a diagnostic, so don't pile on an extra diagnostic here.
Type::Dynamic(DynamicType::Unknown) => {} Type::Dynamic(DynamicType::Unknown) => return,
other_type => { _ => {}
}
if let Some(builder) = self if let Some(builder) = self
.context .context
.report_lint(&INVALID_NEWTYPE, &arguments.args[1]) .report_lint(&INVALID_NEWTYPE, &arguments.args[1])
{ {
let mut diag = builder.into_diagnostic("invalid base for `typing.NewType`"); let mut diag = builder.into_diagnostic("invalid base for `typing.NewType`");
diag.set_primary_message(format!("type `{}`", other_type.display(self.db()))); diag.set_primary_message(format!("type `{}`", inferred.display(self.db())));
if matches!(other_type, Type::ProtocolInstance(_)) { if matches!(inferred, Type::ProtocolInstance(_)) {
diag.info("The base of a `NewType` is not allowed to be a protocol class."); diag.info("The base of a `NewType` is not allowed to be a protocol class.");
} else if matches!(other_type, Type::TypedDict(_)) { } else if matches!(inferred, Type::TypedDict(_)) {
diag.info("The base of a `NewType` is not allowed to be a `TypedDict`."); diag.info("The base of a `NewType` is not allowed to be a `TypedDict`.");
} else { } else {
diag.info( diag.info("The base of a `NewType` must be a class type or another `NewType`.");
"The base of a `NewType` must be a class type or another `NewType`.",
);
}
}
} }
} }
} }

View File

@ -3,7 +3,9 @@ use std::collections::BTreeSet;
use crate::Db; use crate::Db;
use crate::semantic_index::definition::{Definition, DefinitionKind}; use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::types::constraints::ConstraintSet; use crate::types::constraints::ConstraintSet;
use crate::types::{ClassType, KnownClass, Type, UnionType, definition_expression_type, visitor}; use crate::types::{
ClassType, KnownClass, KnownUnion, Type, UnionType, definition_expression_type, visitor,
};
use ruff_db::parsed::parsed_module; use ruff_db::parsed::parsed_module;
use ruff_python_ast as ast; use ruff_python_ast as ast;
@ -84,8 +86,11 @@ impl<'db> NewType<'db> {
// `int | float | complex`. These are allowed because that's what `float` and `complex` // `int | float | complex`. These are allowed because that's what `float` and `complex`
// expand into in type position. We don't currently ask whether the union was implicit // expand into in type position. We don't currently ask whether the union was implicit
// or explicit, so the explicit version is also allowed. // or explicit, so the explicit version is also allowed.
Type::Union(union_type) if union_type.is_int_float(db) => NewTypeBase::Float, Type::Union(union_type) => match union_type.known(db) {
Type::Union(union_type) if union_type.is_int_float_complex(db) => NewTypeBase::Complex, Some(KnownUnion::Float) => NewTypeBase::Float,
Some(KnownUnion::Complex) => NewTypeBase::Complex,
_ => object_fallback,
},
_ => object_fallback, _ => object_fallback,
} }
} }