[ty] support `NewType`s of `float` and `complex` (#21886)

Fixes https://github.com/astral-sh/ty/issues/1818.

crates/ty_python_semantic/resources/mdtest/annotations/new_types.md

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

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::NewTypeInstance(self_newtype),
+                self_newtype.concrete_base_type(db).has_relation_to_impl(
-                Type::NominalInstance(target_nominal_instance),
-            ) => self_newtype.base_class_type(db).has_relation_to_impl(
                     db,
-                target_nominal_instance.class(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
+            // Other than the special cases enumerated above, nominal-instance types are never
-            // newtype-instance types are never subtypes of any other variants
+            // 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::NewTypeInstance(newtype) => newtype.concrete_base_type(db).is_singleton(db),
-                Type::instance(db, newtype.base_class_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::NewTypeInstance(newtype) => newtype.concrete_base_type(db).is_single_valued(db),
-                Type::instance(db, newtype.base_class_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))
+            Type::NewTypeInstance(newtype) => {
-                .try_bool_impl(db, allow_short_circuit, visitor)?,
+                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_type(db, newtype.concrete_base_type(db));
-                visitor.visit_generic_alias_type(db, generic_alias);
-            }
         }
     }
 }
@@ -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<KnownUnion> {
+        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)]

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

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(
+            Type::NewTypeInstance(newtype) => {
-                db,
+                ClassBase::try_from_type(db, newtype.concrete_base_type(db), subclass)
-                Type::instance(db, newtype.base_class_type(db)),
+            }
-                subclass,
-            ),
 
             Type::PropertyInstance(_)
             | Type::BooleanLiteral(_)

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(
+        Type::NewTypeInstance(newtype) => {
-            db,
+            always_true_if(is_instance(&newtype.concrete_base_type(db)))
-            newtype.base_class_type(db),
+        }
-        ))),
 
         Type::BooleanLiteral(..)
         | Type::BytesLiteral(..)

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,
+    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::{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]) {
+        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
+            // `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) => {}
+            Type::Dynamic(DynamicType::Unknown) => return,
-            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())));
+            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.");
-                    } 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`.");
             } 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`.",
-                        );
-                    }
-                }
             }
         }
     }

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

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
+            // There are exactly two union types allowed as bases for NewType: `int | float` and
-            // other newtypes, for example unions. `NewType("Foo", int | str)` is not allowed.
+            // `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
+    // Walk the `NewTypeBase` chain to find the underlying non-newtype `Type`. There might not be
-    // `ClassType` if this `NewType` is cyclical, and we fall back to `object` in that case.
+    // one 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> {
+    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 {
+            match base {
-                return class_type;
+                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
+        // If we get here, there is no `ClassType` (because this newtype is either float/complex or
-        // call `f` at all.
+        // 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<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 +293,10 @@ impl<'db> Iterator for NewTypeBaseIter<'db> {
                     Some(NewTypeBase::NewType(base_newtype))
                 }
             }
+            concrete_base => {
+                self.current = None;
+                Some(concrete_base)
+            }
         }
     }
 }