From 0aeee5009cc0d0cc75d19de6ff7551ee6a47166a Mon Sep 17 00:00:00 2001
From: Jack O'Connor <oconnor663@gmail.com>
Date: Fri, 19 Dec 2025 19:18:23 -0800
Subject: [PATCH] [ty] fix comparisons and arithmetic with `NewType`s of
 `float`

Fixes https://github.com/astral-sh/ty/issues/2077.
---
 .../resources/mdtest/annotations/new_types.md | 12 +++++-
 .../src/types/infer/builder.rs                | 42 +++++++++++++++++--
 2 files changed, 48 insertions(+), 6 deletions(-)

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 ab55503691..d90a41cc9a 100644
--- a/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md
+++ b/crates/ty_python_semantic/resources/mdtest/annotations/new_types.md
@@ -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))
diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs
index bd1a75d9a5..5af6029965 100644
--- a/crates/ty_python_semantic/src/types/infer/builder.rs
+++ b/crates/ty_python_semantic/src/types/infer/builder.rs
@@ -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)),