load suppressions from comment ranges using block range index

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

crates/ruff_linter/src/linter.rs

@@ -404,7 +404,8 @@ pub fn add_noqa_to_path(
     );
 
     // Parse range suppression comments
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
     // Generate diagnostics, ignoring any existing `noqa` directives.
     let diagnostics = check_path(
@@ -470,7 +471,8 @@ pub fn lint_only(
     );
 
     // Parse range suppression comments
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
     // Generate diagnostics.
     let diagnostics = check_path(
@@ -579,7 +581,8 @@ pub fn lint_fix<'a>(
         );
 
         // Parse range suppression comments
-        let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+        let suppressions =
+            Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
         // Generate diagnostics.
         let diagnostics = check_path(
@@ -961,7 +964,8 @@ mod tests {
             &locator,
             &indexer,
         );
-        let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+        let suppressions =
+            Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
         let mut diagnostics = check_path(
             path,
             None,

crates/ruff_linter/src/rules/pyflakes/mod.rs

@@ -957,7 +957,7 @@ mod tests {
             &indexer,
         );
         let suppressions =
-            Suppressions::from_tokens(&settings, locator.contents(), parsed.tokens());
+            Suppressions::from_tokens(&settings, locator.contents(), parsed.tokens(), &indexer);
         let mut messages = check_path(
             Path::new("<filename>"),
             None,

crates/ruff_linter/src/suppression.rs

@@ -2,14 +2,15 @@ use compact_str::CompactString;
 use core::fmt;
 use ruff_db::diagnostic::Diagnostic;
 use ruff_diagnostics::{Edit, Fix};
-use ruff_python_ast::token::{TokenKind, Tokens};
+use ruff_python_ast::token::{Token, TokenKind, Tokens};
 use ruff_python_ast::whitespace::indentation;
-use rustc_hash::FxHashSet;
+use ruff_python_index::Indexer;
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::cell::Cell;
 use std::{error::Error, fmt::Formatter};
 use thiserror::Error;
 
-use ruff_python_trivia::Cursor;
+use ruff_python_trivia::{CommentRanges, Cursor};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize, TextSlice};
 use smallvec::{SmallVec, smallvec};
 
@@ -125,10 +126,17 @@ pub struct Suppressions {
 }
 
 impl Suppressions {
-    pub fn from_tokens(settings: &LinterSettings, source: &str, tokens: &Tokens) -> Suppressions {
+    pub fn from_tokens(
+        settings: &LinterSettings,
+        source: &str,
+        tokens: &Tokens,
+        indexer: &Indexer,
+    ) -> Suppressions {
         if is_range_suppressions_enabled(settings) {
             let builder = SuppressionsBuilder::new(source);
-            builder.load_from_tokens(tokens)
+            builder.load_from_tokens_indexed(tokens, indexer)
+            // let builder = SuppressionsBuilder::new(source);
+            // builder.load_from_tokens(tokens)
         } else {
             Suppressions::default()
         }
@@ -303,7 +311,7 @@ impl Suppressions {
     }
 }
 
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub(crate) struct SuppressionsBuilder<'a> {
     source: &'a str,
 
@@ -312,6 +320,7 @@ pub(crate) struct SuppressionsBuilder<'a> {
     errors: Vec<ParseError>,
 
     pending: Vec<PendingSuppressionComment<'a>>,
+    pending_by_indent: FxHashMap<TextRange, Vec<SuppressionComment>>,
 }
 
 impl<'a> SuppressionsBuilder<'a> {
@@ -322,6 +331,89 @@ impl<'a> SuppressionsBuilder<'a> {
         }
     }
 
+    pub(crate) fn load_from_tokens_indexed(
+        mut self,
+        tokens: &Tokens,
+        indexer: &Indexer,
+    ) -> Suppressions {
+        let global_indent = TextRange::empty(0.into());
+
+        dbg!(indexer.block_ranges());
+        'outer: for comment_range in indexer.comment_ranges() {
+            dbg!(comment_range);
+            dbg!(self.source.slice(&comment_range));
+            let mut parser = SuppressionParser::new(self.source, comment_range);
+            match parser.parse_comment() {
+                Ok(comment) => {
+                    let Some(indent) = indentation(self.source, &comment_range) else {
+                        // trailing suppressions are not supported
+                        self.invalid.push(InvalidSuppression {
+                            kind: InvalidSuppressionKind::Trailing,
+                            comment,
+                        });
+                        continue;
+                    };
+                    let comment_indent = indent.text_len();
+
+                    let token_index = match tokens.binary_search_by_start(comment_range.start()) {
+                        Ok(index) => index,
+                        Err(index) => index,
+                    };
+                    let precedes_dedent = tokens[token_index..]
+                        .iter()
+                        .find(|t| !t.kind().is_trivia())
+                        .is_some_and(|token| matches!(token.kind(), TokenKind::Dedent));
+
+                    let block_ranges = indexer.block_ranges().containing(&comment_range);
+                    dbg!(&block_ranges);
+
+                    // no blocks, global scope
+                    if block_ranges.is_empty() && comment_indent == 0.into() {
+                        self.pending_by_indent
+                            .entry(global_indent)
+                            .or_default()
+                            .push(comment);
+                        continue 'outer;
+                    }
+
+                    for block in indexer
+                        .block_ranges()
+                        .containing(&comment_range)
+                        .iter()
+                        .rev()
+                    {
+                        let block_indent = block.indent.len();
+
+                        if comment_indent == block_indent {
+                            self.pending_by_indent
+                                .entry(block.indent)
+                                .or_default()
+                                .push(comment);
+                            continue 'outer;
+                        } else if comment_indent < block_indent && precedes_dedent {
+                            continue;
+                        }
+                        break;
+                    }
+
+                    // weirdly indented? ¯\_(ツ)_/¯
+                    self.invalid.push(InvalidSuppression {
+                        kind: InvalidSuppressionKind::Indentation,
+                        comment,
+                    });
+                }
+                Err(ParseError {
+                    kind: ParseErrorKind::NotASuppression,
+                    ..
+                }) => {}
+                Err(error) => self.errors.push(error),
+            }
+        }
+
+        dbg!(self);
+
+        Suppressions::default()
+    }
     pub(crate) fn load_from_tokens(mut self, tokens: &Tokens) -> Suppressions {
         let default_indent = "";
         let mut indents: Vec<&str> = vec![];
@@ -642,6 +734,7 @@ mod tests {
 
     use insta::assert_debug_snapshot;
     use itertools::Itertools;
+    use ruff_python_index::Indexer;
     use ruff_python_parser::{Mode, ParseOptions, parse};
     use ruff_text_size::{TextRange, TextSize};
     use similar::DiffableStr;
@@ -1568,10 +1661,12 @@ def bar():
         /// Parse all suppressions and errors in a module for testing
         fn debug(source: &'_ str) -> DebugSuppressions<'_> {
             let parsed = parse(source, ParseOptions::from(Mode::Module)).unwrap();
+            let indexer = Indexer::from_tokens(parsed.tokens(), source);
             let suppressions = Suppressions::from_tokens(
                 &LinterSettings::default().with_preview_mode(),
                 source,
                 parsed.tokens(),
+                &indexer,
             );
             DebugSuppressions {
                 source,

crates/ruff_linter/src/test.rs

@@ -235,7 +235,8 @@ pub(crate) fn test_contents<'a>(
         &locator,
         &indexer,
     );
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
     let messages = check_path(
         path,
         path.parent()
@@ -303,7 +304,7 @@ pub(crate) fn test_contents<'a>(
             );
 
             let suppressions =
-                Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+                Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
             let fixed_messages = check_path(
                 path,
                 None,

crates/ruff_python_index/src/indexer.rs

@@ -4,7 +4,7 @@
 use ruff_python_ast::Stmt;
 use ruff_python_ast::token::{TokenKind, Tokens};
 use ruff_python_trivia::{
-    CommentRanges, has_leading_content, has_trailing_content, is_python_whitespace,
+    BlockRanges, CommentRanges, has_leading_content, has_trailing_content, is_python_whitespace,
 };
 use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -26,6 +26,9 @@ pub struct Indexer {
 
     /// The range of all comments in the source document.
     comment_ranges: CommentRanges,
+
+    /// The ranges of all indent/dedent blocks in the source document.
+    block_ranges: BlockRanges,
 }
 
 impl Indexer {
@@ -36,6 +39,8 @@ impl Indexer {
         let mut multiline_ranges_builder = MultilineRangesBuilder::default();
         let mut continuation_lines = Vec::new();
         let mut comment_ranges = Vec::new();
+        let mut indent_ranges = Vec::new();
+        let mut dedent_ranges = Vec::new();
 
         // Token, end
         let mut prev_end = TextSize::default();
@@ -76,6 +81,12 @@ impl Indexer {
                 TokenKind::Comment => {
                     comment_ranges.push(token.range());
                 }
+                TokenKind::Indent => {
+                    indent_ranges.push(token.range());
+                }
+                TokenKind::Dedent => {
+                    dedent_ranges.push(token.range());
+                }
                 _ => {}
             }
 
@@ -87,6 +98,7 @@ impl Indexer {
             interpolated_string_ranges: interpolated_string_ranges_builder.finish(),
             multiline_ranges: multiline_ranges_builder.finish(),
             comment_ranges: CommentRanges::new(comment_ranges),
+            block_ranges: BlockRanges::new(indent_ranges, dedent_ranges),
         }
     }
 
@@ -95,6 +107,11 @@ impl Indexer {
         &self.comment_ranges
     }
 
+    /// Returns the block indent/dedent ranges.
+    pub const fn block_ranges(&self) -> &BlockRanges {
+        &self.block_ranges
+    }
+
     /// Returns the byte offset ranges of interpolated strings.
     pub const fn interpolated_string_ranges(&self) -> &InterpolatedStringRanges {
         &self.interpolated_string_ranges

crates/ruff_python_trivia/src/block_ranges.rs

@@ -0,0 +1,48 @@
+use ruff_text_size::TextRange;
+
+/// Stores the ranges of indents and dedents sorted by [`TextRange::start`] in increasing order.
+#[derive(Clone, Debug, Default)]
+pub struct BlockRanges {
+    raw: Vec<BlockRange>,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct BlockRange {
+    pub indent: TextRange,
+    pub dedent: TextRange,
+}
+
+impl BlockRanges {
+    pub fn new(indent_ranges: Vec<TextRange>, dedent_ranges: Vec<TextRange>) -> Self {
+        let mut index = 0;
+        let mut stack = Vec::new();
+        let mut blocks = Vec::new();
+
+        for dedent in &dedent_ranges {
+            while index < indent_ranges.len() && indent_ranges[index].end() < dedent.start() {
+                stack.push(indent_ranges[index]);
+                index += 1;
+            }
+
+            if let Some(indent) = stack.pop() {
+                blocks.push(BlockRange {
+                    indent,
+                    dedent: *dedent,
+                });
+            }
+        }
+
+        blocks.sort_by_key(|b| b.indent.start());
+
+        Self { raw: blocks }
+    }
+
+    pub fn containing(&self, range: &TextRange) -> Vec<&BlockRange> {
+        self.raw
+            .iter()
+            .filter(|block| {
+                block.indent.start() <= range.start() && block.dedent.end() > range.end()
+            })
+            .collect()
+    }
+}

crates/ruff_python_trivia/src/lib.rs

@@ -1,3 +1,4 @@
+mod block_ranges;
 mod comment_ranges;
 mod comments;
 mod cursor;
@@ -6,6 +7,7 @@ pub mod textwrap;
 mod tokenizer;
 mod whitespace;
 
+pub use block_ranges::BlockRanges;
 pub use comment_ranges::CommentRanges;
 pub use comments::*;
 pub use cursor::*;

crates/ruff_server/src/lint.rs

@@ -120,8 +120,12 @@ pub(crate) fn check(
     let directives = extract_directives(parsed.tokens(), Flags::all(), &locator, &indexer);
 
     // Parse range suppression comments
-    let suppressions =
-        Suppressions::from_tokens(&settings.linter, locator.contents(), parsed.tokens());
+    let suppressions = Suppressions::from_tokens(
+        &settings.linter,
+        locator.contents(),
+        parsed.tokens(),
+        &indexer,
+    );
 
     // Generate checks.
     let diagnostics = check_path(

crates/ruff_wasm/src/lib.rs

@@ -213,8 +213,12 @@ impl Workspace {
             &indexer,
         );
 
-        let suppressions =
-            Suppressions::from_tokens(&self.settings.linter, locator.contents(), parsed.tokens());
+        let suppressions = Suppressions::from_tokens(
+            &self.settings.linter,
+            locator.contents(),
+            parsed.tokens(),
+            &indexer,
+        );
 
         // Generate checks.
         let diagnostics = check_path(

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

@@ -168,6 +168,56 @@ on top of that:
 Foo = NewType("Foo", 42)
 ```
 
+## `NewType`s in arithmetic and comparison expressions might or might not act as their base
+
+These expressions are valid because `Foo` acts as its base type, `int`:
+
+```py
+from typing import NewType
+
+Foo = NewType("Foo", int)
+
+reveal_type(Foo(42) + 1)  # revealed: int
+reveal_type(1 + Foo(42))  # revealed: int
+reveal_type(Foo(42) + Foo(42))  # revealed: int
+reveal_type(Foo(42) == 42)  # revealed: bool
+reveal_type(42 == Foo(42))  # revealed: bool
+reveal_type(Foo(42) == Foo(42))  # revealed: bool
+```
+
+However, we can't always substitute `int` for `Foo` to evaluate expressions like these. In the
+following cases, only `Foo` itself is valid:
+
+```py
+class Bar:
+    def __add__(self, other: Foo) -> Foo:
+        return other
+
+    def __radd__(self, other: Foo) -> Foo:
+        return other
+
+    def __lt__(self, other: Foo) -> bool:
+        return True
+
+    def __gt__(self, other: Foo) -> bool:
+        return True
+
+    def __contains__(self, key: Foo) -> bool:
+        return True
+
+reveal_type(Foo(42) + Bar())  # revealed: Foo
+reveal_type(Bar() + Foo(42))  # revealed: Foo
+reveal_type(Foo(42) < Bar())  # revealed: bool
+reveal_type(Bar() < Foo(42))  # revealed: bool
+reveal_type(Foo(42) in Bar())  # revealed: bool
+
+42 + Bar()  # error: [unsupported-operator]
+Bar() + 42  # error: [unsupported-operator]
+42 < Bar()  # error: [unsupported-operator]
+Bar() < 42  # error: [unsupported-operator]
+42 in Bar()  # error: [unsupported-operator]
+```
+
 ## `float` and `complex` special cases
 
 `float` and `complex` are subject to a special case in the typing spec, which we currently interpret
@@ -178,6 +228,7 @@ and we accept the unions they expand into.
 
 ```py
 from typing import NewType
+from ty_extensions import static_assert, is_assignable_to
 
 Foo = NewType("Foo", float)
 Foo(3.14)
@@ -186,6 +237,15 @@ Foo("hello")  # error: [invalid-argument-type] "Argument is incorrect: Expected
 
 reveal_type(Foo(3.14).__class__)  # revealed: type[int] | type[float]
 reveal_type(Foo(42).__class__)  # revealed: type[int] | type[float]
+static_assert(is_assignable_to(Foo, float))
+static_assert(is_assignable_to(Foo, int | float))
+static_assert(is_assignable_to(Foo, int | float | None))
+# The assignments above require treating `Foo` as its underlying union type. Each of its members is
+# assignable to the union on the right, so `Foo` is assignable to the union, even though `Foo` as a
+# whole isn't assignable to any one member. However, as in the previous section, we need to be sure
+# that this treatment doesn't break cases like the assignment below, where `Foo` *is* assignable to
+# the union on the right, even though its members *aren't*.
+static_assert(is_assignable_to(Foo, Foo | None))
 
 Bar = NewType("Bar", complex)
 Bar(1 + 2j)
@@ -196,6 +256,11 @@ 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]
+static_assert(is_assignable_to(Bar, complex))
+static_assert(is_assignable_to(Bar, int | float | complex))
+static_assert(is_assignable_to(Bar, int | float | complex | None))
+# See the `Foo | None` case above.
+static_assert(is_assignable_to(Bar, Bar | None))
 ```
 
 We don't currently try to distinguish between an implicit union (e.g. `float`) and the equivalent
@@ -223,6 +288,52 @@ def g(_: Callable[[int | float | complex], Bar]): ...
 g(Bar)
 ```
 
+The arithmetic and comparison test cases in the previous section used a `NewType` of `int`, but
+`NewType`s of `float` and `complex` are more complicated, because their base type is a union, and
+that union needs special handling in binary expressions. In these examples, we we need to lower
+`Foo` to `int | float` and then check each member of that union _individually_, as we would with an
+explicit `Union` on the left side:
+
+```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
+```
+
+But again as above, we can't _always_ lower `Foo` to `int | float`, because there are also binary
+expressions where only `Foo` itself is valid:
+
+```py
+class Bing:
+    def __add__(self, other: Foo) -> Foo:
+        return other
+
+    def __radd__(self, other: Foo) -> Foo:
+        return other
+
+    def __lt__(self, other: Foo) -> bool:
+        return True
+
+    def __gt__(self, other: Foo) -> bool:
+        return True
+
+    def __contains__(self, key: Foo) -> bool:
+        return True
+
+reveal_type(Foo(3.14) + Bing())  # revealed: Foo
+reveal_type(Bing() + Foo(42))  # revealed: Foo
+reveal_type(Foo(3.14) < Bing())  # revealed: bool
+reveal_type(Bing() < Foo(42))  # revealed: bool
+reveal_type(Foo(3.14) in Bing())  # revealed: bool
+
+3.14 + Bing()  # error: [unsupported-operator]
+Bing() + 3.14  # error: [unsupported-operator]
+3.14 < Bing()  # error: [unsupported-operator]
+Bing() < 3.14  # error: [unsupported-operator]
+3.14 in Bing()  # error: [unsupported-operator]
+```
+
 ## A `NewType` definition must be a simple variable assignment
 
 ```py

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -10426,6 +10426,41 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 op,
             ),
 
+            // `try_call_bin_op` works for almost all `NewType`s, but not for `NewType`s of `float`
+            // and `complex`, where the concrete base type is a union. In that case it turns out
+            // the `self` types of the dunder methods in typeshed don't match, because they don't
+            // get the same `int | float` and `int | float | complex` special treatment that the
+            // positional arguments get. In those cases we need to explicitly delegate to the base
+            // type, so that it hits the `Type::Union` branches above.
+            (Type::NewTypeInstance(newtype), rhs, _) => {
+                Type::try_call_bin_op(self.db(), left_ty, op, right_ty)
+                    .map(|outcome| outcome.return_type(self.db()))
+                    .ok()
+                    .or_else(|| {
+                        self.infer_binary_expression_type(
+                            node,
+                            emitted_division_by_zero_diagnostic,
+                            newtype.concrete_base_type(self.db()),
+                            rhs,
+                            op,
+                        )
+                    })
+            }
+            (lhs, Type::NewTypeInstance(newtype), _) => {
+                Type::try_call_bin_op(self.db(), left_ty, op, right_ty)
+                    .map(|outcome| outcome.return_type(self.db()))
+                    .ok()
+                    .or_else(|| {
+                        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(_)), _, _)
@@ -10762,8 +10797,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 | Type::TypeVar(_)
                 | Type::TypeIs(_)
                 | Type::TypeGuard(_)
-                | Type::TypedDict(_)
-                | Type::NewTypeInstance(_),
+                | Type::TypedDict(_),
                 Type::FunctionLiteral(_)
                 | Type::BooleanLiteral(_)
                 | Type::Callable(..)
@@ -10793,8 +10827,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 | Type::TypeVar(_)
                 | Type::TypeIs(_)
                 | Type::TypeGuard(_)
-                | 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()))
@@ -11228,6 +11261,39 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 )
             })),
 
+            // `try_dunder` works for almost all `NewType`s, but not for `NewType`s of `float` and
+            // `complex`, where the concrete base type is a union. In that case it turns out the
+            // `self` types of the dunder methods in typeshed don't match, because they don't get
+            // the same `int | float` and `int | float | complex` special treatment that the
+            // positional arguments get. In those cases we need to explicitly delegate to the base
+            // type, so that it hits the `Type::Union` branches above.
+            (Type::NewTypeInstance(newtype), right) => Some(
+                try_dunder(self, MemberLookupPolicy::default()).or_else(|_| {
+                    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(
+                try_dunder(self, MemberLookupPolicy::default()).or_else(|_| {
+                    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)),

crates/ty_python_semantic/src/types/relation.rs

@@ -654,6 +654,79 @@ impl<'db> Type<'db> {
             // `Never` is the bottom type, the empty set.
             (_, Type::Never) => ConstraintSet::from(false),
 
+            (Type::NewTypeInstance(self_newtype), Type::NewTypeInstance(target_newtype)) => {
+                self_newtype.has_relation_to_impl(db, target_newtype)
+            }
+            // In the special cases of `NewType`s of `float` or `complex`, the concrete base type
+            // can be a union (`int | float` or `int | float | complex`). For that reason,
+            // `NewType` assignability to a union needs to consider two different cases. It could
+            // be that we need to treat the `NewType` as the underlying union it's assignable to,
+            // for example:
+            //
+            // ```py
+            // Foo = NewType("Foo", float)
+            // static_assert(is_assignable_to(Foo, float | None))
+            // ```
+            //
+            // The right side there is equivalent to `int | float | None`, but `Foo` as a whole
+            // isn't assignable to any of those three types. However, `Foo`s concrete base type is
+            // `int | float`, which is assignable, because union members on the left side get
+            // checked individually. On the other hand, we need to be careful not to break the
+            // following case, where `int | float` is *not* assignable to the right side:
+            //
+            // ```py
+            // static_assert(is_assignable_to(Foo, Foo | None))
+            // ```
+            //
+            // To handle both cases, we have to check that *either* `Foo` as a whole is assignable
+            // (or subtypeable etc.) *or* that its concrete base type is. Note that this match arm
+            // needs to take precedence over the `Type::Union` arms immediately below.
+            (Type::NewTypeInstance(self_newtype), Type::Union(union)) => {
+                // First the normal "assign to union" case, unfortunately duplicated from below.
+                union
+                    .elements(db)
+                    .iter()
+                    .when_any(db, |&elem_ty| {
+                        self.has_relation_to_impl(
+                            db,
+                            elem_ty,
+                            inferable,
+                            relation,
+                            relation_visitor,
+                            disjointness_visitor,
+                        )
+                    })
+                    // Failing that, if the concrete base type is a union, try delegating to that.
+                    // Otherwise, this would be equivalent to what we just checked, and we
+                    // shouldn't waste time checking it twice.
+                    .or(db, || {
+                        let concrete_base = self_newtype.concrete_base_type(db);
+                        if matches!(concrete_base, Type::Union(_)) {
+                            concrete_base.has_relation_to_impl(
+                                db,
+                                target,
+                                inferable,
+                                relation,
+                                relation_visitor,
+                                disjointness_visitor,
+                            )
+                        } else {
+                            ConstraintSet::from(false)
+                        }
+                    })
+            }
+            // All other `NewType` assignments fall back to the concrete base type.
+            (Type::NewTypeInstance(self_newtype), _) => {
+                self_newtype.concrete_base_type(db).has_relation_to_impl(
+                    db,
+                    target,
+                    inferable,
+                    relation,
+                    relation_visitor,
+                    disjointness_visitor,
+                )
+            }
+
             (Type::Union(union), _) => union.elements(db).iter().when_all(db, |&elem_ty| {
                 elem_ty.has_relation_to_impl(
                     db,
@@ -1305,21 +1378,6 @@ impl<'db> Type<'db> {
                 })
             }
 
-            (Type::NewTypeInstance(self_newtype), Type::NewTypeInstance(target_newtype)) => {
-                self_newtype.has_relation_to_impl(db, target_newtype)
-            }
-
-            (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(
                     db,