diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md b/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md index 39a88cff49..ab55503691 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md @@ -146,9 +146,10 @@ Foo = NewType(name, int) reveal_type(Foo) # revealed: ``` -## The second argument must be a class type or another newtype +## The base must be a class type or another newtype -Other typing constructs like `Union` are not allowed. +Other typing constructs like `Union` are not _generally_ allowed. (However, see the next section for +a couple special cases.) ```py from typing_extensions import NewType @@ -167,6 +168,61 @@ on top of that: Foo = NewType("Foo", 42) ``` +## `float` and `complex` special cases + +`float` and `complex` are subject to a special case in the typing spec, which we currently interpret +to mean that `float` in type position is `int | float`, and `complex` in type position is +`int | float | complex`. This is awkward for `NewType`, because as we just tested above, unions +aren't generally valid `NewType` bases. However, `float` and `complex` _are_ valid `NewType` bases, +and we accept the unions they expand into. + +```py +from typing import NewType + +Foo = NewType("Foo", float) +Foo(3.14) +Foo(42) +Foo("hello") # error: [invalid-argument-type] "Argument is incorrect: Expected `int | float`, found `Literal["hello"]`" + +reveal_type(Foo(3.14).__class__) # revealed: type[int] | type[float] +reveal_type(Foo(42).__class__) # revealed: type[int] | type[float] + +Bar = NewType("Bar", complex) +Bar(1 + 2j) +Bar(3.14) +Bar(42) +Bar("goodbye") # error: [invalid-argument-type] + +reveal_type(Bar(1 + 2j).__class__) # revealed: type[int] | type[float] | type[complex] +reveal_type(Bar(3.14).__class__) # revealed: type[int] | type[float] | type[complex] +reveal_type(Bar(42).__class__) # revealed: type[int] | type[float] | type[complex] +``` + +We don't currently try to distinguish between an implicit union (e.g. `float`) and the equivalent +explicit union (e.g. `int | float`), so these two explicit unions are also allowed. But again, most +unions are not allowed: + +```py +Baz = NewType("Baz", int | float) +Baz = NewType("Baz", int | float | complex) +Baz = NewType("Baz", int | str) # error: [invalid-newtype] "invalid base for `typing.NewType`" +``` + +Similarly, a `NewType` of `float` or `complex` is valid as a `Callable` of the corresponding union +type: + +```py +from collections.abc import Callable + +def f(_: Callable[[int | float], Foo]): ... + +f(Foo) + +def g(_: Callable[[int | float | complex], Bar]): ... + +g(Bar) +``` + ## A `NewType` definition must be a simple variable assignment ```py @@ -179,7 +235,7 @@ N: NewType = NewType("N", int) # error: [invalid-newtype] "A `NewType` definiti Cyclic newtypes are kind of silly, but it's possible for the user to express them, and it's important that we don't go into infinite recursive loops and crash with a stack overflow. In fact, -this is *why* base type evaluation is deferred; otherwise Salsa itself would crash. +this is _why_ base type evaluation is deferred; otherwise Salsa itself would crash. ```py from typing_extensions import NewType, reveal_type, cast diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index f531fb604e..cfbafef322 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -1785,7 +1785,7 @@ impl<'db> Type<'db> { Type::GenericAlias(alias) => Some(ClassType::Generic(alias).into_callable(db)), Type::NewTypeInstance(newtype) => { - Type::instance(db, newtype.base_class_type(db)).try_upcast_to_callable(db) + newtype.concrete_base_type(db).try_upcast_to_callable(db) } // TODO: This is unsound so in future we can consider an opt-in option to disable it. @@ -2906,17 +2906,16 @@ impl<'db> Type<'db> { self_newtype.has_relation_to_impl(db, target_newtype) } - ( - Type::NewTypeInstance(self_newtype), - Type::NominalInstance(target_nominal_instance), - ) => self_newtype.base_class_type(db).has_relation_to_impl( - db, - target_nominal_instance.class(db), - inferable, - relation, - relation_visitor, - disjointness_visitor, - ), + (Type::NewTypeInstance(self_newtype), _) => { + self_newtype.concrete_base_type(db).has_relation_to_impl( + db, + target, + inferable, + relation, + relation_visitor, + disjointness_visitor, + ) + } (Type::PropertyInstance(_), _) => { KnownClass::Property.to_instance(db).has_relation_to_impl( @@ -2937,10 +2936,9 @@ impl<'db> Type<'db> { disjointness_visitor, ), - // Other than the special cases enumerated above, nominal-instance types, and - // newtype-instance types are never subtypes of any other variants + // Other than the special cases enumerated above, nominal-instance types are never + // subtypes of any other variants (Type::NominalInstance(_), _) => ConstraintSet::from(false), - (Type::NewTypeInstance(_), _) => ConstraintSet::from(false), } } @@ -3869,7 +3867,7 @@ impl<'db> Type<'db> { left.is_disjoint_from_impl(db, right) } (Type::NewTypeInstance(newtype), other) | (other, Type::NewTypeInstance(newtype)) => { - Type::instance(db, newtype.base_class_type(db)).is_disjoint_from_impl( + newtype.concrete_base_type(db).is_disjoint_from_impl( db, other, inferable, @@ -4100,9 +4098,7 @@ impl<'db> Type<'db> { Type::TypeIs(type_is) => type_is.is_bound(db), Type::TypedDict(_) => false, Type::TypeAlias(alias) => alias.value_type(db).is_singleton(db), - Type::NewTypeInstance(newtype) => { - Type::instance(db, newtype.base_class_type(db)).is_singleton(db) - } + Type::NewTypeInstance(newtype) => newtype.concrete_base_type(db).is_singleton(db), } } @@ -4153,9 +4149,7 @@ impl<'db> Type<'db> { } Type::NominalInstance(instance) => instance.is_single_valued(db), - Type::NewTypeInstance(newtype) => { - Type::instance(db, newtype.base_class_type(db)).is_single_valued(db) - } + Type::NewTypeInstance(newtype) => newtype.concrete_base_type(db).is_single_valued(db), Type::BoundSuper(_) => { // At runtime two super instances never compare equal, even if their arguments are identical. @@ -4407,7 +4401,9 @@ impl<'db> Type<'db> { Type::Dynamic(_) | Type::Never => Place::bound(self).into(), Type::NominalInstance(instance) => instance.class(db).instance_member(db, name), - Type::NewTypeInstance(newtype) => newtype.base_class_type(db).instance_member(db, name), + Type::NewTypeInstance(newtype) => { + newtype.concrete_base_type(db).instance_member(db, name) + } Type::ProtocolInstance(protocol) => protocol.instance_member(db, name), @@ -5523,8 +5519,11 @@ impl<'db> Type<'db> { .value_type(db) .try_bool_impl(db, allow_short_circuit, visitor) })?, - Type::NewTypeInstance(newtype) => Type::instance(db, newtype.base_class_type(db)) - .try_bool_impl(db, allow_short_circuit, visitor)?, + Type::NewTypeInstance(newtype) => { + newtype + .concrete_base_type(db) + .try_bool_impl(db, allow_short_circuit, visitor)? + } }; Ok(truthiness) @@ -6492,7 +6491,7 @@ impl<'db> Type<'db> { match ty { Type::NominalInstance(nominal) => nominal.tuple_spec(db), - Type::NewTypeInstance(newtype) => non_async_special_case(db, Type::instance(db, newtype.base_class_type(db))), + Type::NewTypeInstance(newtype) => non_async_special_case(db, newtype.concrete_base_type(db)), Type::GenericAlias(alias) if alias.origin(db).is_tuple(db) => { Some(Cow::Owned(TupleSpec::homogeneous(todo_type!( "*tuple[] annotations" @@ -7605,7 +7604,7 @@ impl<'db> Type<'db> { ), }, Type::TypeAlias(alias) => alias.value_type(db).to_meta_type(db), - Type::NewTypeInstance(newtype) => Type::from(newtype.base_class_type(db)), + Type::NewTypeInstance(newtype) => newtype.concrete_base_type(db).to_meta_type(db), } } @@ -8811,9 +8810,7 @@ fn walk_known_instance_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>( visitor.visit_callable_type(db, callable); } KnownInstanceType::NewType(newtype) => { - if let ClassType::Generic(generic_alias) = newtype.base_class_type(db) { - visitor.visit_generic_alias_type(db, generic_alias); - } + visitor.visit_type(db, newtype.concrete_base_type(db)); } } } @@ -13918,6 +13915,39 @@ impl<'db> UnionType<'db> { ConstraintSet::from(sorted_self == other.normalized(db)) } + + /// Identify some specific unions of known classes, currently the ones that `float` and + /// `complex` expand into in type position. + pub(crate) fn known(self, db: &'db dyn Db) -> Option { + let mut has_int = false; + let mut has_float = false; + let mut has_complex = false; + for element in self.elements(db) { + 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, + 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, + } + } +} + +#[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)] diff --git a/crates/ty_python_semantic/src/types/bound_super.rs b/crates/ty_python_semantic/src/types/bound_super.rs index 442ae0d0b9..cea3bc3e54 100644 --- a/crates/ty_python_semantic/src/types/bound_super.rs +++ b/crates/ty_python_semantic/src/types/bound_super.rs @@ -429,7 +429,7 @@ impl<'db> BoundSuperType<'db> { ); } Type::NewTypeInstance(newtype) => { - return delegate_to(Type::instance(db, newtype.base_class_type(db))); + return delegate_to(newtype.concrete_base_type(db)); } Type::Callable(callable) if callable.is_function_like(db) => { return delegate_to(KnownClass::FunctionType.to_instance(db)); diff --git a/crates/ty_python_semantic/src/types/class_base.rs b/crates/ty_python_semantic/src/types/class_base.rs index c5ce573d3b..26b490fa3b 100644 --- a/crates/ty_python_semantic/src/types/class_base.rs +++ b/crates/ty_python_semantic/src/types/class_base.rs @@ -152,11 +152,9 @@ impl<'db> ClassBase<'db> { Type::TypeAlias(alias) => Self::try_from_type(db, alias.value_type(db), subclass), - Type::NewTypeInstance(newtype) => ClassBase::try_from_type( - db, - Type::instance(db, newtype.base_class_type(db)), - subclass, - ), + Type::NewTypeInstance(newtype) => { + ClassBase::try_from_type(db, newtype.concrete_base_type(db), subclass) + } Type::PropertyInstance(_) | Type::BooleanLiteral(_) diff --git a/crates/ty_python_semantic/src/types/function.rs b/crates/ty_python_semantic/src/types/function.rs index 6b61316a69..e727e6663b 100644 --- a/crates/ty_python_semantic/src/types/function.rs +++ b/crates/ty_python_semantic/src/types/function.rs @@ -1247,10 +1247,9 @@ fn is_instance_truthiness<'db>( Type::NominalInstance(..) => always_true_if(is_instance(&ty)), - Type::NewTypeInstance(newtype) => always_true_if(is_instance(&Type::instance( - db, - newtype.base_class_type(db), - ))), + Type::NewTypeInstance(newtype) => { + always_true_if(is_instance(&newtype.concrete_base_type(db))) + } Type::BooleanLiteral(..) | Type::BytesLiteral(..) diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 00141d4ca8..e4a488dc7e 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -104,13 +104,13 @@ use crate::types::visitor::any_over_type; use crate::types::{ BoundTypeVarInstance, CallDunderError, CallableBinding, CallableType, CallableTypeKind, ClassLiteral, ClassType, DataclassParams, DynamicType, InternedType, IntersectionBuilder, - IntersectionType, KnownClass, KnownInstanceType, LintDiagnosticGuard, MemberLookupPolicy, - MetaclassCandidate, PEP695TypeAliasType, ParamSpecAttrKind, Parameter, ParameterForm, - Parameters, Signature, SpecialFormType, SubclassOfType, TrackedConstraintSet, Truthiness, Type, - TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers, TypeVarBoundOrConstraints, - TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarIdentity, - TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder, UnionType, - UnionTypeInstance, binding_type, infer_scope_types, todo_type, + IntersectionType, KnownClass, KnownInstanceType, KnownUnion, LintDiagnosticGuard, + MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, ParamSpecAttrKind, Parameter, + ParameterForm, Parameters, Signature, SpecialFormType, SubclassOfType, TrackedConstraintSet, + Truthiness, Type, TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers, + TypeVarBoundOrConstraints, TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, + TypeVarIdentity, TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder, + UnionType, UnionTypeInstance, binding_type, infer_scope_types, todo_type, }; use crate::types::{CallableTypes, overrides}; use crate::types::{ClassBase, add_inferred_python_version_hint_to_diagnostic}; @@ -5629,28 +5629,35 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // Infer the deferred base type of a NewType. fn infer_newtype_assignment_deferred(&mut self, arguments: &ast::Arguments) { - match self.infer_type_expression(&arguments.args[1]) { - Type::NominalInstance(_) | Type::NewTypeInstance(_) => {} + let inferred = self.infer_type_expression(&arguments.args[1]); + match inferred { + Type::NominalInstance(_) | Type::NewTypeInstance(_) => return, + // 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` + // expand into in type position. We don't currently ask whether the union was implicit + // or explicit, so the explicit version is also allowed. + Type::Union(union_ty) => { + if let Some(KnownUnion::Float | KnownUnion::Complex) = union_ty.known(self.db()) { + return; + } + } // `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. - Type::Dynamic(DynamicType::Unknown) => {} - other_type => { - if let Some(builder) = self - .context - .report_lint(&INVALID_NEWTYPE, &arguments.args[1]) - { - let mut diag = builder.into_diagnostic("invalid base for `typing.NewType`"); - diag.set_primary_message(format!("type `{}`", other_type.display(self.db()))); - if matches!(other_type, Type::ProtocolInstance(_)) { - diag.info("The base of a `NewType` is not allowed to be a protocol class."); - } else if matches!(other_type, Type::TypedDict(_)) { - diag.info("The base of a `NewType` is not allowed to be a `TypedDict`."); - } else { - diag.info( - "The base of a `NewType` must be a class type or another `NewType`.", - ); - } - } + Type::Dynamic(DynamicType::Unknown) => return, + _ => {} + } + if let Some(builder) = self + .context + .report_lint(&INVALID_NEWTYPE, &arguments.args[1]) + { + let mut diag = builder.into_diagnostic("invalid base for `typing.NewType`"); + diag.set_primary_message(format!("type `{}`", inferred.display(self.db()))); + if matches!(inferred, Type::ProtocolInstance(_)) { + diag.info("The base of a `NewType` is not allowed to be a protocol class."); + } else if matches!(inferred, Type::TypedDict(_)) { + diag.info("The base of a `NewType` is not allowed to be a `TypedDict`."); + } else { + diag.info("The base of a `NewType` must be a class type or another `NewType`."); } } } diff --git a/crates/ty_python_semantic/src/types/list_members.rs b/crates/ty_python_semantic/src/types/list_members.rs index a93438a596..4e4a32c294 100644 --- a/crates/ty_python_semantic/src/types/list_members.rs +++ b/crates/ty_python_semantic/src/types/list_members.rs @@ -187,7 +187,7 @@ impl<'db> AllMembers<'db> { } Type::NewTypeInstance(newtype) => { - self.extend_with_type(db, Type::instance(db, newtype.base_class_type(db))); + self.extend_with_type(db, newtype.concrete_base_type(db)); } Type::ClassLiteral(class_literal) if class_literal.is_typed_dict(db) => { diff --git a/crates/ty_python_semantic/src/types/newtype.rs b/crates/ty_python_semantic/src/types/newtype.rs index 84a6e18f50..cc6f2cff69 100644 --- a/crates/ty_python_semantic/src/types/newtype.rs +++ b/crates/ty_python_semantic/src/types/newtype.rs @@ -3,7 +3,9 @@ use std::collections::BTreeSet; use crate::Db; use crate::semantic_index::definition::{Definition, DefinitionKind}; use crate::types::constraints::ConstraintSet; -use crate::types::{ClassType, Type, definition_expression_type, visitor}; +use crate::types::{ + ClassType, KnownClass, KnownUnion, Type, UnionType, definition_expression_type, visitor, +}; use ruff_db::parsed::parsed_module; use ruff_python_ast as ast; @@ -80,8 +82,15 @@ impl<'db> NewType<'db> { NewTypeBase::ClassType(nominal_instance_type.class(db)) } Type::NewTypeInstance(newtype) => NewTypeBase::NewType(newtype), - // This branch includes bases that are other typing constructs besides classes and - // other newtypes, for example unions. `NewType("Foo", int | str)` is not allowed. + // 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` + // expand into in type position. We don't currently ask whether the union was implicit + // or explicit, so the explicit version is also allowed. + Type::Union(union_type) => match union_type.known(db) { + Some(KnownUnion::Float) => NewTypeBase::Float, + Some(KnownUnion::Complex) => NewTypeBase::Complex, + _ => object_fallback, + }, _ => object_fallback, } } @@ -94,15 +103,16 @@ impl<'db> NewType<'db> { } } - // Walk the `NewTypeBase` chain to find the underlying `ClassType`. There might not be a - // `ClassType` if this `NewType` is cyclical, and we fall back to `object` in that case. - pub fn base_class_type(self, db: &'db dyn Db) -> ClassType<'db> { + // Walk the `NewTypeBase` chain to find the underlying non-newtype `Type`. There might not be + // one if this `NewType` is cyclical, and we fall back to `object` in that case. + pub fn concrete_base_type(self, db: &'db dyn Db) -> Type<'db> { for base in self.iter_bases(db) { - if let NewTypeBase::ClassType(class_type) = base { - return class_type; + match base { + NewTypeBase::NewType(_) => continue, + concrete => return concrete.instance_type(db), } } - ClassType::object(db) + Type::object() } pub(crate) fn is_equivalent_to_impl(self, db: &'db dyn Db, other: Self) -> bool { @@ -179,10 +189,14 @@ impl<'db> NewType<'db> { Some(mapped_base), )); } + // Mapping base class types is used for normalization and applying type mappings, + // neither of which have any effect on `float` or `complex` (which are already + // fully normalized and non-generic), so we don't need to bother calling `f`. + NewTypeBase::Float | NewTypeBase::Complex => {} } } - // If we get here, there is no `ClassType` (because this newtype is cyclic), and we don't - // call `f` at all. + // If we get here, there is no `ClassType` (because this newtype is either float/complex or + // cyclic), and we don't call `f` at all. Some(self) } @@ -209,6 +223,12 @@ pub(crate) fn walk_newtype_instance_type<'db, V: visitor::TypeVisitor<'db> + ?Si pub enum NewTypeBase<'db> { ClassType(ClassType<'db>), NewType(NewType<'db>), + // `float` and `complex` are special-cased in type position, where they refer to `int | float` + // and `int | float | complex` respectively. As an extension of that special case, we allow + // them in `NewType` bases, even though unions and other typing constructs normally aren't + // allowed. + Float, + Complex, } impl<'db> NewTypeBase<'db> { @@ -216,6 +236,21 @@ impl<'db> NewTypeBase<'db> { match self { NewTypeBase::ClassType(class_type) => Type::instance(db, class_type), NewTypeBase::NewType(newtype) => Type::NewTypeInstance(newtype), + NewTypeBase::Float => UnionType::from_elements( + db, + [ + KnownClass::Int.to_instance(db), + KnownClass::Float.to_instance(db), + ], + ), + NewTypeBase::Complex => UnionType::from_elements( + db, + [ + KnownClass::Int.to_instance(db), + KnownClass::Float.to_instance(db), + KnownClass::Complex.to_instance(db), + ], + ), } } } @@ -246,10 +281,6 @@ impl<'db> Iterator for NewTypeBaseIter<'db> { fn next(&mut self) -> Option { let current = self.current?; match current.base(self.db) { - NewTypeBase::ClassType(base_class_type) => { - self.current = None; - Some(NewTypeBase::ClassType(base_class_type)) - } NewTypeBase::NewType(base_newtype) => { // Doing the insertion only in this branch avoids allocating in the common case. self.seen_before.insert(current); @@ -262,6 +293,10 @@ impl<'db> Iterator for NewTypeBaseIter<'db> { Some(NewTypeBase::NewType(base_newtype)) } } + concrete_base => { + self.current = None; + Some(concrete_base) + } } } }