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 55afa55a0b..c50b13b32a 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -10058,6 +10058,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(_)), _, _) @@ -10393,8 +10409,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(..) @@ -10423,8 +10438,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())) @@ -10858,6 +10872,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)),