[ty] fix comparisons and arithmetic with NewTypes of float

Fixes https://github.com/astral-sh/ty/issues/2077.
This commit is contained in:
Jack O'Connor 2025-12-19 19:18:23 -08:00
parent 674d3902c6
commit 0aeee5009c
2 changed files with 48 additions and 6 deletions

View file

@ -223,6 +223,15 @@ def g(_: Callable[[int | float | complex], Bar]): ...
g(Bar)
```
Arithmetic also works:
```py
reveal_type(Foo(3.14) < Foo(42)) # revealed: bool
reveal_type(Foo(3.14) == Foo(42)) # revealed: bool
reveal_type(Foo(3.14) + Foo(42)) # revealed: int | float
reveal_type(Foo(3.14) / Foo(42)) # revealed: int | float
```
## A `NewType` definition must be a simple variable assignment
```py
@ -300,8 +309,7 @@ A = NewType("A", EllipsisType)
static_assert(is_singleton(A))
static_assert(is_single_valued(A))
reveal_type(type(A(...)) is EllipsisType) # revealed: Literal[True]
# TODO: This should be `Literal[True]` also.
reveal_type(A(...) is ...) # revealed: bool
reveal_type(A(...) is ...) # revealed: Literal[True]
B = NewType("B", int)
static_assert(not is_singleton(B))

View file

@ -9993,6 +9993,22 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
op,
),
(Type::NewTypeInstance(newtype), rhs, _) => self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
newtype.concrete_base_type(self.db()),
rhs,
op,
),
(lhs, Type::NewTypeInstance(newtype), _) => self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
lhs,
newtype.concrete_base_type(self.db()),
op,
),
// Non-todo Anys take precedence over Todos (as if we fix this `Todo` in the future,
// the result would then become Any or Unknown, respectively).
(div @ Type::Dynamic(DynamicType::Divergent(_)), _, _)
@ -10328,8 +10344,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::BoundSuper(_)
| Type::TypeVar(_)
| Type::TypeIs(_)
| Type::TypedDict(_)
| Type::NewTypeInstance(_),
| Type::TypedDict(_),
Type::FunctionLiteral(_)
| Type::BooleanLiteral(_)
| Type::Callable(..)
@ -10358,8 +10373,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
| Type::BoundSuper(_)
| Type::TypeVar(_)
| Type::TypeIs(_)
| Type::TypedDict(_)
| Type::NewTypeInstance(_),
| Type::TypedDict(_),
op,
) => Type::try_call_bin_op(self.db(), left_ty, op, right_ty)
.map(|outcome| outcome.return_type(self.db()))
@ -10793,6 +10807,26 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
)
})),
(Type::NewTypeInstance(newtype), right) => Some(
visitor.visit((left, op, right), || { self.infer_binary_type_comparison(
newtype.concrete_base_type(self.db()),
op,
right,
range,
visitor,
)
})),
(left, Type::NewTypeInstance(newtype)) => Some(
visitor.visit((left, op, right), || { self.infer_binary_type_comparison(
left,
op,
newtype.concrete_base_type(self.db()),
range,
visitor,
)
})),
(Type::IntLiteral(n), Type::IntLiteral(m)) => Some(match op {
ast::CmpOp::Eq => Ok(Type::BooleanLiteral(n == m)),
ast::CmpOp::NotEq => Ok(Type::BooleanLiteral(n != m)),