Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[ty] support `NewType`s of `float` and `complex` 

https://github.com/astral-sh/ty/issues/1818
This commit is contained in:
Jack O'Connor 2025-12-09 18:47:58 -08:00
parent c8851ecf70
commit 26891484c7
8 changed files with 191 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
crates/ty_python_semantic/resources/mdtest/annotations/new_types.md
View File

@ -146,9 +146,10 @@ Foo = NewType(name, int)
reveal_type(Foo) # revealed: <NewType pseudo-class 'Foo'>
```
## 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

105
crates/ty_python_semantic/src/types.rs
View File

@ -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,54 @@ impl<'db> UnionType<'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
/// `complex` expands into in type position.
pub(crate) fn is_int_float_complex(self, db: &'db dyn Db) -> bool {
let elements = self.elements(db);
if elements.len() != 3 {
return false;
}
let mut has_int = false;
let mut has_float = false;
let mut has_complex = 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,
KnownClass::Complex => has_complex = true,
_ => {}
}
}
}
has_int && has_float && has_complex
}
}
#[salsa::interned(debug, heap_size=IntersectionType::heap_size)]

2
crates/ty_python_semantic/src/types/bound_super.rs
View File

@ -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));

8
crates/ty_python_semantic/src/types/class_base.rs
View File

@ -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(_)

7
crates/ty_python_semantic/src/types/function.rs
View File

@ -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(..)

7
crates/ty_python_semantic/src/types/infer/builder.rs
View File

@ -5631,6 +5631,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
fn infer_newtype_assignment_deferred(&mut self, arguments: &ast::Arguments) {
match self.infer_type_expression(&arguments.args[1]) {
Type::NominalInstance(_) | Type::NewTypeInstance(_) => {}
// 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 union_ty.is_int_float(self.db()) || union_ty.is_int_float_complex(self.db()) => {
}
// `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) => {}

2
crates/ty_python_semantic/src/types/list_members.rs
View File

@ -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) => {

57
crates/ty_python_semantic/src/types/newtype.rs
View File

@ -3,7 +3,7 @@ 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, Type, UnionType, definition_expression_type, visitor};
use ruff_db::parsed::parsed_module;
use ruff_python_ast as ast;
@ -80,8 +80,12 @@ 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) if union_type.is_int_float(db) => NewTypeBase::Float,
Type::Union(union_type) if union_type.is_int_float_complex(db) => NewTypeBase::Complex,
_ => object_fallback,
}
}
@ -94,15 +98,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 +184,11 @@ impl<'db> NewType<'db> {
Some(mapped_base),
));
}
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 +215,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 +228,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 +273,6 @@ impl<'db> Iterator for NewTypeBaseIter<'db> {
fn next(&mut self) -> Option<Self::Item> {
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 +285,10 @@ impl<'db> Iterator for NewTypeBaseIter<'db> {
Some(NewTypeBase::NewType(base_newtype))
}
}
concrete_base => {
self.current = None;
Some(concrete_base)
}
}
}
}