Remove SuppressionsIter

I really misunderstood this in
https://github.com/astral-sh/ruff/pull/15664#issuecomment-2613079710

.gitignore

@@ -30,7 +30,7 @@ tracing-flamechart.svg
 tracing-flamegraph.svg
 
 # insta
-.rs.pending-snap
+*.rs.pending-snap
 
 
 ###

crates/red_knot/src/args.rs

@@ -0,0 +1,190 @@
+use crate::logging::Verbosity;
+use crate::python_version::PythonVersion;
+use clap::{ArgAction, ArgMatches, Error, Parser};
+use red_knot_project::metadata::options::{EnvironmentOptions, Options};
+use red_knot_project::metadata::value::{RangedValue, RelativePathBuf};
+use red_knot_python_semantic::lint;
+use ruff_db::system::SystemPathBuf;
+
+#[derive(Debug, Parser)]
+#[command(
+    author,
+    name = "red-knot",
+    about = "An extremely fast Python type checker."
+)]
+#[command(version)]
+pub(crate) struct Args {
+    #[command(subcommand)]
+    pub(crate) command: Command,
+}
+
+#[derive(Debug, clap::Subcommand)]
+pub(crate) enum Command {
+    /// Check a project for type errors.
+    Check(CheckCommand),
+
+    /// Start the language server
+    Server,
+}
+
+#[derive(Debug, Parser)]
+pub(crate) struct CheckCommand {
+    /// Run the command within the given project directory.
+    ///
+    /// All `pyproject.toml` files will be discovered by walking up the directory tree from the given project directory,
+    /// as will the project's virtual environment (`.venv`) unless the `venv-path` option is set.
+    ///
+    /// Other command-line arguments (such as relative paths) will be resolved relative to the current working directory.
+    #[arg(long, value_name = "PROJECT")]
+    pub(crate) project: Option<SystemPathBuf>,
+
+    /// Path to the virtual environment the project uses.
+    ///
+    /// If provided, red-knot will use the `site-packages` directory of this virtual environment
+    /// to resolve type information for the project's third-party dependencies.
+    #[arg(long, value_name = "PATH")]
+    pub(crate) venv_path: Option<SystemPathBuf>,
+
+    /// Custom directory to use for stdlib typeshed stubs.
+    #[arg(long, value_name = "PATH", alias = "custom-typeshed-dir")]
+    pub(crate) typeshed: Option<SystemPathBuf>,
+
+    /// Additional path to use as a module-resolution source (can be passed multiple times).
+    #[arg(long, value_name = "PATH")]
+    pub(crate) extra_search_path: Option<Vec<SystemPathBuf>>,
+
+    /// Python version to assume when resolving types.
+    #[arg(long, value_name = "VERSION", alias = "target-version")]
+    pub(crate) python_version: Option<PythonVersion>,
+
+    #[clap(flatten)]
+    pub(crate) verbosity: Verbosity,
+
+    #[clap(flatten)]
+    pub(crate) rules: RulesArg,
+
+    /// Run in watch mode by re-running whenever files change.
+    #[arg(long, short = 'W')]
+    pub(crate) watch: bool,
+}
+
+impl CheckCommand {
+    pub(crate) fn into_options(self) -> Options {
+        let rules = if self.rules.is_empty() {
+            None
+        } else {
+            Some(
+                self.rules
+                    .into_iter()
+                    .map(|(rule, level)| (RangedValue::cli(rule), RangedValue::cli(level)))
+                    .collect(),
+            )
+        };
+
+        Options {
+            environment: Some(EnvironmentOptions {
+                python_version: self
+                    .python_version
+                    .map(|version| RangedValue::cli(version.into())),
+                venv_path: self.venv_path.map(RelativePathBuf::cli),
+                typeshed: self.typeshed.map(RelativePathBuf::cli),
+                extra_paths: self.extra_search_path.map(|extra_search_paths| {
+                    extra_search_paths
+                        .into_iter()
+                        .map(RelativePathBuf::cli)
+                        .collect()
+                }),
+                ..EnvironmentOptions::default()
+            }),
+            rules,
+            ..Default::default()
+        }
+    }
+}
+
+/// A list of rules to enable or disable with a given severity.
+///
+/// This type is used to parse the `--error`, `--warn`, and `--ignore` arguments
+/// while preserving the order in which they were specified (arguments last override previous severities).
+#[derive(Debug)]
+pub(crate) struct RulesArg(Vec<(String, lint::Level)>);
+
+impl RulesArg {
+    fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    fn into_iter(self) -> impl Iterator<Item = (String, lint::Level)> {
+        self.0.into_iter()
+    }
+}
+
+impl clap::FromArgMatches for RulesArg {
+    fn from_arg_matches(matches: &ArgMatches) -> Result<Self, Error> {
+        let mut rules = Vec::new();
+
+        for (level, arg_id) in [
+            (lint::Level::Ignore, "ignore"),
+            (lint::Level::Warn, "warn"),
+            (lint::Level::Error, "error"),
+        ] {
+            let indices = matches.indices_of(arg_id).into_iter().flatten();
+            let levels = matches.get_many::<String>(arg_id).into_iter().flatten();
+            rules.extend(
+                indices
+                    .zip(levels)
+                    .map(|(index, rule)| (index, rule, level)),
+            );
+        }
+
+        // Sort by their index so that values specified later override earlier ones.
+        rules.sort_by_key(|(index, _, _)| *index);
+
+        Ok(Self(
+            rules
+                .into_iter()
+                .map(|(_, rule, level)| (rule.to_owned(), level))
+                .collect(),
+        ))
+    }
+
+    fn update_from_arg_matches(&mut self, matches: &ArgMatches) -> Result<(), Error> {
+        self.0 = Self::from_arg_matches(matches)?.0;
+        Ok(())
+    }
+}
+
+impl clap::Args for RulesArg {
+    fn augment_args(cmd: clap::Command) -> clap::Command {
+        const HELP_HEADING: &str = "Enabling / disabling rules";
+
+        cmd.arg(
+            clap::Arg::new("error")
+                .long("error")
+                .action(ArgAction::Append)
+                .help("Treat the given rule as having severity 'error'. Can be specified multiple times.")
+                .value_name("RULE")
+                .help_heading(HELP_HEADING),
+        )
+        .arg(
+            clap::Arg::new("warn")
+                .long("warn")
+                .action(ArgAction::Append)
+                .help("Treat the given rule as having severity 'warn'. Can be specified multiple times.")
+                .value_name("RULE")
+                .help_heading(HELP_HEADING),
+        )
+        .arg(
+            clap::Arg::new("ignore")
+                .long("ignore")
+                .action(ArgAction::Append)
+                .help("Disables the rule. Can be specified multiple times.")
+                .value_name("RULE")
+                .help_heading(HELP_HEADING),
+        )
+    }
+
+    fn augment_args_for_update(cmd: clap::Command) -> clap::Command {
+        Self::augment_args(cmd)
+    }
+}

crates/red_knot/src/main.rs

@@ -1,13 +1,13 @@
 use std::process::{ExitCode, Termination};
 use std::sync::Mutex;
 
+use crate::args::{Args, CheckCommand, Command};
+use crate::logging::setup_tracing;
 use anyhow::{anyhow, Context};
 use clap::Parser;
 use colored::Colorize;
 use crossbeam::channel as crossbeam_channel;
-use python_version::PythonVersion;
-use red_knot_project::metadata::options::{EnvironmentOptions, Options};
-use red_knot_project::metadata::value::{RangedValue, RelativePathBuf};
+use red_knot_project::metadata::options::Options;
 use red_knot_project::watch;
 use red_knot_project::watch::ProjectWatcher;
 use red_knot_project::{ProjectDatabase, ProjectMetadata};
@@ -16,87 +16,11 @@ use ruff_db::diagnostic::Diagnostic;
 use ruff_db::system::{OsSystem, System, SystemPath, SystemPathBuf};
 use salsa::plumbing::ZalsaDatabase;
 
-use crate::logging::{setup_tracing, Verbosity};
-
+mod args;
 mod logging;
 mod python_version;
 mod verbosity;
 
-#[derive(Debug, Parser)]
-#[command(
-    author,
-    name = "red-knot",
-    about = "An extremely fast Python type checker."
-)]
-#[command(version)]
-struct Args {
-    #[command(subcommand)]
-    pub(crate) command: Option<Command>,
-
-    /// Run the command within the given project directory.
-    ///
-    /// All `pyproject.toml` files will be discovered by walking up the directory tree from the given project directory,
-    /// as will the project's virtual environment (`.venv`) unless the `venv-path` option is set.
-    ///
-    /// Other command-line arguments (such as relative paths) will be resolved relative to the current working directory.
-    #[arg(long, value_name = "PROJECT")]
-    project: Option<SystemPathBuf>,
-
-    /// Path to the virtual environment the project uses.
-    ///
-    /// If provided, red-knot will use the `site-packages` directory of this virtual environment
-    /// to resolve type information for the project's third-party dependencies.
-    #[arg(long, value_name = "PATH")]
-    venv_path: Option<SystemPathBuf>,
-
-    /// Custom directory to use for stdlib typeshed stubs.
-    #[arg(long, value_name = "PATH", alias = "custom-typeshed-dir")]
-    typeshed: Option<SystemPathBuf>,
-
-    /// Additional path to use as a module-resolution source (can be passed multiple times).
-    #[arg(long, value_name = "PATH")]
-    extra_search_path: Option<Vec<SystemPathBuf>>,
-
-    /// Python version to assume when resolving types.
-    #[arg(long, value_name = "VERSION", alias = "target-version")]
-    python_version: Option<PythonVersion>,
-
-    #[clap(flatten)]
-    verbosity: Verbosity,
-
-    /// Run in watch mode by re-running whenever files change.
-    #[arg(long, short = 'W')]
-    watch: bool,
-}
-
-impl Args {
-    fn to_options(&self) -> Options {
-        Options {
-            environment: Some(EnvironmentOptions {
-                python_version: self
-                    .python_version
-                    .map(|version| RangedValue::cli(version.into())),
-                venv_path: self.venv_path.as_ref().map(RelativePathBuf::cli),
-                typeshed: self.typeshed.as_ref().map(RelativePathBuf::cli),
-                extra_paths: self.extra_search_path.as_ref().map(|extra_search_paths| {
-                    extra_search_paths
-                        .iter()
-                        .map(RelativePathBuf::cli)
-                        .collect()
-                }),
-                ..EnvironmentOptions::default()
-            }),
-            ..Default::default()
-        }
-    }
-}
-
-#[derive(Debug, clap::Subcommand)]
-pub enum Command {
-    /// Start the language server
-    Server,
-}
-
 #[allow(clippy::print_stdout, clippy::unnecessary_wraps, clippy::print_stderr)]
 pub fn main() -> ExitStatus {
     run().unwrap_or_else(|error| {
@@ -122,10 +46,13 @@ pub fn main() -> ExitStatus {
 fn run() -> anyhow::Result<ExitStatus> {
     let args = Args::parse_from(std::env::args());
 
-    if matches!(args.command, Some(Command::Server)) {
-        return run_server().map(|()| ExitStatus::Success);
+    match args.command {
+        Command::Server => run_server().map(|()| ExitStatus::Success),
+        Command::Check(check_args) => run_check(check_args),
     }
+}
 
+fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
     let verbosity = args.verbosity.level();
     countme::enable(verbosity.is_trace());
     let _guard = setup_tracing(verbosity)?;
@@ -156,7 +83,8 @@ fn run() -> anyhow::Result<ExitStatus> {
         .unwrap_or_else(|| cli_base_path.clone());
 
     let system = OsSystem::new(cwd);
-    let cli_options = args.to_options();
+    let watch = args.watch;
+    let cli_options = args.into_options();
     let mut workspace_metadata = ProjectMetadata::discover(system.current_directory(), &system)?;
     workspace_metadata.apply_cli_options(cli_options.clone());
 
@@ -174,7 +102,7 @@ fn run() -> anyhow::Result<ExitStatus> {
         }
     })?;
 
-    let exit_status = if args.watch {
+    let exit_status = if watch {
         main_loop.watch(&mut db)?
     } else {
         main_loop.run(&mut db)

crates/red_knot/tests/cli.rs

@@ -1,5 +1,5 @@
 use anyhow::Context;
-use insta::Settings;
+use insta::internals::SettingsBindDropGuard;
 use insta_cmd::{assert_cmd_snapshot, get_cargo_bin};
 use std::path::{Path, PathBuf};
 use std::process::Command;
@@ -28,24 +28,22 @@ fn config_override() -> anyhow::Result<()> {
         ),
     ])?;
 
-    case.insta_settings().bind(|| {
-        assert_cmd_snapshot!(case.command(), @r"
-            success: false
-            exit_code: 1
-            ----- stdout -----
-            error[lint:unresolved-attribute] <temp_dir>/test.py:5:7 Type `<module 'sys'>` has no attribute `last_exc`
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        error[lint:unresolved-attribute] <temp_dir>/test.py:5:7 Type `<module 'sys'>` has no attribute `last_exc`
 
-            ----- stderr -----
-        ");
+        ----- stderr -----
+    ");
 
-        assert_cmd_snapshot!(case.command().arg("--python-version").arg("3.12"), @r"
-            success: true
-            exit_code: 0
-            ----- stdout -----
+    assert_cmd_snapshot!(case.command().arg("--python-version").arg("3.12"), @r"
+        success: true
+        exit_code: 0
+        ----- stdout -----
 
-            ----- stderr -----
-        ");
-    });
+        ----- stderr -----
+    ");
 
     Ok(())
 }
@@ -92,25 +90,23 @@ fn cli_arguments_are_relative_to_the_current_directory() -> anyhow::Result<()> {
         ),
     ])?;
 
-    case.insta_settings().bind(|| {
-        // Make sure that the CLI fails when the `libs` directory is not in the search path.
-        assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r#"
-            success: false
-            exit_code: 1
-            ----- stdout -----
-            error[lint:unresolved-import] <temp_dir>/child/test.py:2:1 Cannot resolve import `utils`
+    // Make sure that the CLI fails when the `libs` directory is not in the search path.
+    assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r#"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        error[lint:unresolved-import] <temp_dir>/child/test.py:2:1 Cannot resolve import `utils`
 
-            ----- stderr -----
-        "#);
+        ----- stderr -----
+    "#);
 
-        assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")).arg("--extra-search-path").arg("../libs"), @r"
-            success: true
-            exit_code: 0
-            ----- stdout -----
+    assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")).arg("--extra-search-path").arg("../libs"), @r"
+        success: true
+        exit_code: 0
+        ----- stdout -----
 
-            ----- stderr -----
-        ");
-    });
+        ----- stderr -----
+    ");
 
     Ok(())
 }
@@ -156,22 +152,20 @@ fn paths_in_configuration_files_are_relative_to_the_project_root() -> anyhow::Re
         ),
     ])?;
 
-    case.insta_settings().bind(|| {
-        assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r"
-            success: true
-            exit_code: 0
-            ----- stdout -----
+    assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r"
+        success: true
+        exit_code: 0
+        ----- stdout -----
 
-            ----- stderr -----
-        ");
-    });
+        ----- stderr -----
+    ");
 
     Ok(())
 }
 
 /// The rule severity can be changed in the configuration file
 #[test]
-fn rule_severity() -> anyhow::Result<()> {
+fn configuration_rule_severity() -> anyhow::Result<()> {
     let case = TestCase::with_file(
         "test.py",
         r#"
@@ -184,41 +178,146 @@ fn rule_severity() -> anyhow::Result<()> {
             "#,
     )?;
 
-    case.insta_settings().bind(|| {
-        // Assert that there's a possibly unresolved reference diagnostic
-        // and that division-by-zero has a severity of error by default.
-        assert_cmd_snapshot!(case.command(), @r"
-            success: false
-            exit_code: 1
-            ----- stdout -----
-            error[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
-            warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:7:7 Name `x` used when possibly not defined
+    // Assert that there's a possibly unresolved reference diagnostic
+    // and that division-by-zero has a severity of error by default.
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        error[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:7:7 Name `x` used when possibly not defined
 
-            ----- stderr -----
-        ");
+        ----- stderr -----
+    ");
 
-        case.write_file("pyproject.toml", r#"
-            [tool.knot.rules]
-            division-by-zero = "warn" # demote to warn
-            possibly-unresolved-reference = "ignore"
-        "#)?;
+    case.write_file(
+        "pyproject.toml",
+        r#"
+        [tool.knot.rules]
+        division-by-zero = "warn" # demote to warn
+        possibly-unresolved-reference = "ignore"
+    "#,
+    )?;
 
-        assert_cmd_snapshot!(case.command(), @r"
-            success: false
-            exit_code: 1
-            ----- stdout -----
-            warning[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        warning[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
 
-            ----- stderr -----
-        ");
+        ----- stderr -----
+    ");
 
-        Ok(())
-    })
+    Ok(())
 }
 
-/// Red Knot warns about unknown rules
+/// The rule severity can be changed using `--ignore`, `--warn`, and `--error`
 #[test]
-fn unknown_rules() -> anyhow::Result<()> {
+fn cli_rule_severity() -> anyhow::Result<()> {
+    let case = TestCase::with_file(
+        "test.py",
+        r#"
+        import does_not_exit
+
+        y = 4 / 0
+
+        for a in range(0, y):
+            x = a
+
+        print(x)  # possibly-unresolved-reference
+        "#,
+    )?;
+
+    // Assert that there's a possibly unresolved reference diagnostic
+    // and that division-by-zero has a severity of error by default.
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        error[lint:unresolved-import] <temp_dir>/test.py:2:8 Cannot resolve import `does_not_exit`
+        error[lint:division-by-zero] <temp_dir>/test.py:4:5 Cannot divide object of type `Literal[4]` by zero
+        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:9:7 Name `x` used when possibly not defined
+
+        ----- stderr -----
+    ");
+
+    assert_cmd_snapshot!(
+        case
+            .command()
+            .arg("--ignore")
+            .arg("possibly-unresolved-reference")
+            .arg("--warn")
+            .arg("division-by-zero")
+            .arg("--warn")
+            .arg("unresolved-import"),
+        @r"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    warning[lint:unresolved-import] <temp_dir>/test.py:2:8 Cannot resolve import `does_not_exit`
+    warning[lint:division-by-zero] <temp_dir>/test.py:4:5 Cannot divide object of type `Literal[4]` by zero
+
+    ----- stderr -----
+    "
+    );
+
+    Ok(())
+}
+
+/// The rule severity can be changed using `--ignore`, `--warn`, and `--error` and
+/// values specified last override previous severities.
+#[test]
+fn cli_rule_severity_precedence() -> anyhow::Result<()> {
+    let case = TestCase::with_file(
+        "test.py",
+        r#"
+        y = 4 / 0
+
+        for a in range(0, y):
+            x = a
+
+        print(x)  # possibly-unresolved-reference
+        "#,
+    )?;
+
+    // Assert that there's a possibly unresolved reference diagnostic
+    // and that division-by-zero has a severity of error by default.
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        error[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:7:7 Name `x` used when possibly not defined
+
+        ----- stderr -----
+    ");
+
+    assert_cmd_snapshot!(
+        case
+            .command()
+            .arg("--error")
+            .arg("possibly-unresolved-reference")
+            .arg("--warn")
+            .arg("division-by-zero")
+            // Override the error severity with warning
+            .arg("--ignore")
+            .arg("possibly-unresolved-reference"),
+        @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        warning[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+
+        ----- stderr -----
+        "
+    );
+
+    Ok(())
+}
+
+/// Red Knot warns about unknown rules specified in a configuration file
+#[test]
+fn configuration_unknown_rules() -> anyhow::Result<()> {
     let case = TestCase::with_files([
         (
             "pyproject.toml",
@@ -230,22 +329,38 @@ fn unknown_rules() -> anyhow::Result<()> {
         ("test.py", "print(10)"),
     ])?;
 
-    case.insta_settings().bind(|| {
-        assert_cmd_snapshot!(case.command(), @r"
-            success: false
-            exit_code: 1
-            ----- stdout -----
-            warning[unknown-rule] <temp_dir>/pyproject.toml:3:1 Unknown lint rule `division-by-zer`
+    assert_cmd_snapshot!(case.command(), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        warning[unknown-rule] <temp_dir>/pyproject.toml:3:1 Unknown lint rule `division-by-zer`
 
-            ----- stderr -----
-        ");
-    });
+        ----- stderr -----
+    ");
+
+    Ok(())
+}
+
+/// Red Knot warns about unknown rules specified in a CLI argument
+#[test]
+fn cli_unknown_rules() -> anyhow::Result<()> {
+    let case = TestCase::with_file("test.py", "print(10)")?;
+
+    assert_cmd_snapshot!(case.command().arg("--ignore").arg("division-by-zer"), @r"
+        success: false
+        exit_code: 1
+        ----- stdout -----
+        warning[unknown-rule] Unknown lint rule `division-by-zer`
+
+        ----- stderr -----
+    ");
 
     Ok(())
 }
 
 struct TestCase {
     _temp_dir: TempDir,
+    _settings_scope: SettingsBindDropGuard,
     project_dir: PathBuf,
 }
 
@@ -260,9 +375,16 @@ impl TestCase {
             .canonicalize()
             .context("Failed to canonicalize project path")?;
 
+        let mut settings = insta::Settings::clone_current();
+        settings.add_filter(&tempdir_filter(&project_dir), "<temp_dir>/");
+        settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
+
+        let settings_scope = settings.bind_to_scope();
+
         Ok(Self {
             project_dir,
             _temp_dir: temp_dir,
+            _settings_scope: settings_scope,
         })
     }
 
@@ -307,17 +429,9 @@ impl TestCase {
         &self.project_dir
     }
 
-    // Returns the insta filters to escape paths in snapshots
-    fn insta_settings(&self) -> Settings {
-        let mut settings = insta::Settings::clone_current();
-        settings.add_filter(&tempdir_filter(&self.project_dir), "<temp_dir>/");
-        settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
-        settings
-    }
-
     fn command(&self) -> Command {
         let mut command = Command::new(get_cargo_bin("red_knot"));
-        command.current_dir(&self.project_dir);
+        command.current_dir(&self.project_dir).arg("check");
         command
     }
 }

crates/red_knot/tests/file_watching.rs

@@ -47,7 +47,7 @@ impl TestCase {
         #[track_caller]
         fn panic_with_formatted_events(events: Vec<ChangeEvent>) -> Vec<ChangeEvent> {
             panic!(
-                "Didn't observe expected change:\n{}",
+                "Didn't observe the expected event. The following events occurred:\n{}",
                 events
                     .into_iter()
                     .map(|event| format!("  - {event:?}"))

crates/red_knot_project/src/metadata/options.rs

@@ -149,6 +149,16 @@ impl Options {
                             format!("Unknown lint rule `{rule_name}`"),
                             Severity::Warning,
                         ),
+                        GetLintError::PrefixedWithCategory { suggestion, .. } => {
+                            OptionDiagnostic::new(
+                                DiagnosticId::UnknownRule,
+                                format!(
+                                    "Unknown lint rule `{rule_name}`. Did you mean `{suggestion}`?"
+                                ),
+                                Severity::Warning,
+                            )
+                        }
+
                         GetLintError::Removed(_) => OptionDiagnostic::new(
                             DiagnosticId::UnknownRule,
                             format!("Unknown lint rule `{rule_name}`"),
@@ -206,6 +216,16 @@ pub struct Rules {
     inner: FxHashMap<RangedValue<String>, RangedValue<Level>>,
 }
 
+impl FromIterator<(RangedValue<String>, RangedValue<Level>)> for Rules {
+    fn from_iter<T: IntoIterator<Item = (RangedValue<String>, RangedValue<Level>)>>(
+        iter: T,
+    ) -> Self {
+        Self {
+            inner: iter.into_iter().collect(),
+        }
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum KnotTomlError {
     #[error(transparent)]

crates/red_knot_python_semantic/resources/mdtest/suppressions/knot_ignore.md

@@ -180,3 +180,11 @@ a = 4 / 0  # error: [division-by-zero]
 # error: [unknown-rule] "Unknown rule `is-equal-14`"
 a = 10 + 4  # knot: ignore[is-equal-14]
 ```
+
+## Code with `lint:` prefix
+
+```py
+# error:[unknown-rule] "Unknown rule `lint:division-by-zero`. Did you mean `division-by-zero`?"
+# error: [division-by-zero]
+a = 10 / 0  # knot: ignore[lint:division-by-zero]
+```

crates/red_knot_python_semantic/src/lint.rs

@@ -1,5 +1,5 @@
 use itertools::Itertools;
-use ruff_db::diagnostic::{LintName, Severity};
+use ruff_db::diagnostic::{DiagnosticId, LintName, Severity};
 use rustc_hash::FxHashMap;
 use std::hash::Hasher;
 use thiserror::Error;
@@ -345,7 +345,18 @@ impl LintRegistry {
                 }
             }
             Some(LintEntry::Removed(lint)) => Err(GetLintError::Removed(lint.name())),
-            None => Err(GetLintError::Unknown(code.to_string())),
+            None => {
+                if let Some(without_prefix) = DiagnosticId::strip_category(code) {
+                    if let Some(entry) = self.by_name.get(without_prefix) {
+                        return Err(GetLintError::PrefixedWithCategory {
+                            prefixed: code.to_string(),
+                            suggestion: entry.id().name.to_string(),
+                        });
+                    }
+                }
+
+                Err(GetLintError::Unknown(code.to_string()))
+            }
         }
     }
 
@@ -382,12 +393,20 @@ impl LintRegistry {
 #[derive(Error, Debug, Clone, PartialEq, Eq)]
 pub enum GetLintError {
     /// The name maps to this removed lint.
-    #[error("lint {0} has been removed")]
+    #[error("lint `{0}` has been removed")]
     Removed(LintName),
 
     /// No lint with the given name is known.
-    #[error("unknown lint {0}")]
+    #[error("unknown lint `{0}`")]
     Unknown(String),
+
+    /// The name uses the full qualified diagnostic id `lint:<rule>` instead of just `rule`.
+    /// The String is the name without the `lint:` category prefix.
+    #[error("unknown lint `{prefixed}`. Did you mean `{suggestion}`?")]
+    PrefixedWithCategory {
+        prefixed: String,
+        suggestion: String,
+    },
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
@@ -399,6 +418,16 @@ pub enum LintEntry {
     Alias(LintId),
 }
 
+impl LintEntry {
+    fn id(self) -> LintId {
+        match self {
+            LintEntry::Lint(id) => id,
+            LintEntry::Removed(id) => id,
+            LintEntry::Alias(id) => id,
+        }
+    }
+}
+
 impl From<&'static LintMetadata> for LintEntry {
     fn from(metadata: &'static LintMetadata) -> Self {
         if metadata.status.is_removed() {

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -1,4 +1,3 @@
-use std::iter::FusedIterator;
 use std::sync::Arc;
 
 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
@@ -6,7 +5,7 @@ use salsa::plumbing::AsId;
 
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
-use ruff_index::{IndexSlice, IndexVec};
+use ruff_index::IndexVec;
 
 use crate::module_name::ModuleName;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
@@ -17,7 +16,6 @@ use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolTable,
 };
-use crate::semantic_index::use_def::UseDefMap;
 use crate::Db;
 
 pub mod ast_ids;
@@ -29,8 +27,7 @@ pub mod symbol;
 mod use_def;
 
 pub(crate) use self::use_def::{
-    BindingWithConstraints, BindingWithConstraintsIterator, DeclarationWithConstraint,
-    DeclarationsIterator, ScopedVisibilityConstraintId,
+    BindingWithConstraints, DeclarationWithConstraint, ScopedVisibilityConstraintId, UseDefMap,
 };
 
 type SymbolMap = hashbrown::HashMap<ScopedSymbolId, (), FxBuildHasher>;
@@ -206,20 +203,43 @@ impl<'db> SemanticIndex<'db> {
 
     /// Returns an iterator over the descendent scopes of `scope`.
     #[allow(unused)]
-    pub(crate) fn descendent_scopes(&self, scope: FileScopeId) -> DescendentsIter {
-        DescendentsIter::new(self, scope)
+    pub(crate) fn descendent_scopes(
+        &self,
+        scope_id: FileScopeId,
+    ) -> impl Iterator<Item = (FileScopeId, &Scope)> + '_ {
+        let scope = &self.scopes[scope_id];
+        let scopes = &self.scopes[scope.descendents.clone()];
+        let mut next_id = scope_id + 1;
+        scopes.iter().map(move |descendent| {
+            let result = (next_id, descendent);
+            next_id = next_id + 1;
+            result
+        })
     }
 
     /// Returns an iterator over the direct child scopes of `scope`.
     #[allow(unused)]
-    pub(crate) fn child_scopes(&self, scope: FileScopeId) -> ChildrenIter {
-        ChildrenIter::new(self, scope)
+    pub(crate) fn child_scopes(
+        &self,
+        scope_id: FileScopeId,
+    ) -> impl Iterator<Item = (FileScopeId, &Scope)> + '_ {
+        self.descendent_scopes(scope_id)
+            .filter(move |(_, scope)| scope.parent == Some(scope_id))
     }
 
     /// Returns an iterator over all ancestors of `scope`, starting with `scope` itself.
     #[allow(unused)]
-    pub(crate) fn ancestor_scopes(&self, scope: FileScopeId) -> AncestorsIter {
-        AncestorsIter::new(self, scope)
+    pub(crate) fn ancestor_scopes(
+        &self,
+        scope_id: FileScopeId,
+    ) -> impl Iterator<Item = (FileScopeId, &Scope)> + '_ {
+        let mut next_id = Some(scope_id);
+        std::iter::from_fn(move || {
+            let current_id = next_id?;
+            let current = &self.scopes[current_id];
+            next_id = current.parent;
+            Some((current_id, current))
+        })
     }
 
     /// Returns the [`Definition`] salsa ingredient for `definition_key`.
@@ -266,98 +286,6 @@ impl<'db> SemanticIndex<'db> {
     }
 }
 
-pub struct AncestorsIter<'a> {
-    scopes: &'a IndexSlice<FileScopeId, Scope>,
-    next_id: Option<FileScopeId>,
-}
-
-impl<'a> AncestorsIter<'a> {
-    fn new(module_symbol_table: &'a SemanticIndex, start: FileScopeId) -> Self {
-        Self {
-            scopes: &module_symbol_table.scopes,
-            next_id: Some(start),
-        }
-    }
-}
-
-impl<'a> Iterator for AncestorsIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let current_id = self.next_id?;
-        let current = &self.scopes[current_id];
-        self.next_id = current.parent;
-
-        Some((current_id, current))
-    }
-}
-
-impl FusedIterator for AncestorsIter<'_> {}
-
-pub struct DescendentsIter<'a> {
-    next_id: FileScopeId,
-    descendents: std::slice::Iter<'a, Scope>,
-}
-
-impl<'a> DescendentsIter<'a> {
-    fn new(symbol_table: &'a SemanticIndex, scope_id: FileScopeId) -> Self {
-        let scope = &symbol_table.scopes[scope_id];
-        let scopes = &symbol_table.scopes[scope.descendents.clone()];
-
-        Self {
-            next_id: scope_id + 1,
-            descendents: scopes.iter(),
-        }
-    }
-}
-
-impl<'a> Iterator for DescendentsIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let descendent = self.descendents.next()?;
-        let id = self.next_id;
-        self.next_id = self.next_id + 1;
-
-        Some((id, descendent))
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        self.descendents.size_hint()
-    }
-}
-
-impl FusedIterator for DescendentsIter<'_> {}
-
-impl ExactSizeIterator for DescendentsIter<'_> {}
-
-pub struct ChildrenIter<'a> {
-    parent: FileScopeId,
-    descendents: DescendentsIter<'a>,
-}
-
-impl<'a> ChildrenIter<'a> {
-    fn new(module_symbol_table: &'a SemanticIndex, parent: FileScopeId) -> Self {
-        let descendents = DescendentsIter::new(module_symbol_table, parent);
-
-        Self {
-            parent,
-            descendents,
-        }
-    }
-}
-
-impl<'a> Iterator for ChildrenIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.descendents
-            .find(|(_, scope)| scope.parent == Some(self.parent))
-    }
-}
-
-impl FusedIterator for ChildrenIter<'_> {}
-
 #[cfg(test)]
 mod tests {
     use ruff_db::files::{system_path_to_file, File};

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -255,11 +255,8 @@
 //! snapshot, and merging a snapshot into the current state. The logic using these methods lives in
 //! [`SemanticIndexBuilder`](crate::semantic_index::builder::SemanticIndexBuilder), e.g. where it
 //! visits a `StmtIf` node.
-use self::symbol_state::{
-    BindingIdWithConstraintsIterator, ConstraintIdIterator, DeclarationIdIterator,
-    ScopedDefinitionId, SymbolBindings, SymbolDeclarations, SymbolState,
-};
 pub(crate) use self::symbol_state::{ScopedConstraintId, ScopedVisibilityConstraintId};
+use self::symbol_state::{ScopedDefinitionId, SymbolBindings, SymbolDeclarations, SymbolState};
 use crate::semantic_index::ast_ids::ScopedUseId;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::ScopedSymbolId;
@@ -286,7 +283,7 @@ pub(crate) struct UseDefMap<'db> {
     all_constraints: AllConstraints<'db>,
 
     /// Array of [`VisibilityConstraint`]s in this scope.
-    visibility_constraints: VisibilityConstraints<'db>,
+    pub(crate) visibility_constraints: VisibilityConstraints<'db>,
 
     /// [`SymbolBindings`] reaching a [`ScopedUseId`].
     bindings_by_use: IndexVec<ScopedUseId, SymbolBindings>,
@@ -313,21 +310,24 @@ impl<'db> UseDefMap<'db> {
     pub(crate) fn bindings_at_use(
         &self,
         use_id: ScopedUseId,
-    ) -> BindingWithConstraintsIterator<'_, 'db> {
+    ) -> impl Iterator<Item = BindingWithConstraints<'db, impl Iterator<Item = Constraint<'db>> + '_>> + '_
+    {
         self.bindings_iterator(&self.bindings_by_use[use_id])
     }
 
     pub(crate) fn public_bindings(
         &self,
         symbol: ScopedSymbolId,
-    ) -> BindingWithConstraintsIterator<'_, 'db> {
+    ) -> impl Iterator<Item = BindingWithConstraints<'db, impl Iterator<Item = Constraint<'db>> + '_>> + '_
+    {
         self.bindings_iterator(self.public_symbols[symbol].bindings())
     }
 
     pub(crate) fn bindings_at_declaration(
         &self,
         declaration: Definition<'db>,
-    ) -> BindingWithConstraintsIterator<'_, 'db> {
+    ) -> impl Iterator<Item = BindingWithConstraints<'db, impl Iterator<Item = Constraint<'db>> + '_>> + '_
+    {
         if let SymbolDefinitions::Bindings(bindings) = &self.definitions_by_definition[&declaration]
         {
             self.bindings_iterator(bindings)
@@ -336,10 +336,10 @@ impl<'db> UseDefMap<'db> {
         }
     }
 
-    pub(crate) fn declarations_at_binding<'map>(
-        &'map self,
+    pub(crate) fn declarations_at_binding(
+        &self,
         binding: Definition<'db>,
-    ) -> DeclarationsIterator<'map, 'db> {
+    ) -> impl Iterator<Item = DeclarationWithConstraint<'db>> + '_ {
         if let SymbolDefinitions::Declarations(declarations) =
             &self.definitions_by_definition[&binding]
         {
@@ -349,10 +349,10 @@ impl<'db> UseDefMap<'db> {
         }
     }
 
-    pub(crate) fn public_declarations<'map>(
-        &'map self,
+    pub(crate) fn public_declarations(
+        &self,
         symbol: ScopedSymbolId,
-    ) -> DeclarationsIterator<'map, 'db> {
+    ) -> impl Iterator<Item = DeclarationWithConstraint<'db>> + '_ {
         let declarations = self.public_symbols[symbol].declarations();
         self.declarations_iterator(declarations)
     }
@@ -360,24 +360,35 @@ impl<'db> UseDefMap<'db> {
     fn bindings_iterator<'map>(
         &'map self,
         bindings: &'map SymbolBindings,
-    ) -> BindingWithConstraintsIterator<'map, 'db> {
-        BindingWithConstraintsIterator {
-            all_definitions: &self.all_definitions,
-            all_constraints: &self.all_constraints,
-            visibility_constraints: &self.visibility_constraints,
-            inner: bindings.iter(),
-        }
+    ) -> impl Iterator<
+        Item = BindingWithConstraints<'db, impl Iterator<Item = Constraint<'db>> + 'map>,
+    > + 'map {
+        bindings
+            .iter()
+            .map(|binding_id_with_constraints| BindingWithConstraints {
+                binding: self.all_definitions[binding_id_with_constraints.definition],
+                constraints: binding_id_with_constraints
+                    .constraint_ids
+                    .map(|constraint_id| self.all_constraints[constraint_id]),
+                visibility_constraint: binding_id_with_constraints.visibility_constraint,
+            })
     }
 
     fn declarations_iterator<'map>(
         &'map self,
         declarations: &'map SymbolDeclarations,
-    ) -> DeclarationsIterator<'map, 'db> {
-        DeclarationsIterator {
-            all_definitions: &self.all_definitions,
-            visibility_constraints: &self.visibility_constraints,
-            inner: declarations.iter(),
-        }
+    ) -> impl Iterator<Item = DeclarationWithConstraint<'db>> + 'map {
+        declarations.iter().map(
+            move |DeclarationIdWithConstraint {
+                      definition,
+                      visibility_constraint,
+                  }| {
+                DeclarationWithConstraint {
+                    declaration: self.all_definitions[definition],
+                    visibility_constraint,
+                }
+            },
+        )
     }
 }
 
@@ -388,89 +399,17 @@ enum SymbolDefinitions {
     Declarations(SymbolDeclarations),
 }
 
-#[derive(Debug)]
-pub(crate) struct BindingWithConstraintsIterator<'map, 'db> {
-    all_definitions: &'map IndexVec<ScopedDefinitionId, Option<Definition<'db>>>,
-    all_constraints: &'map AllConstraints<'db>,
-    pub(crate) visibility_constraints: &'map VisibilityConstraints<'db>,
-    inner: BindingIdWithConstraintsIterator<'map>,
-}
-
-impl<'map, 'db> Iterator for BindingWithConstraintsIterator<'map, 'db> {
-    type Item = BindingWithConstraints<'map, 'db>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let all_constraints = self.all_constraints;
-
-        self.inner
-            .next()
-            .map(|binding_id_with_constraints| BindingWithConstraints {
-                binding: self.all_definitions[binding_id_with_constraints.definition],
-                constraints: ConstraintsIterator {
-                    all_constraints,
-                    constraint_ids: binding_id_with_constraints.constraint_ids,
-                },
-                visibility_constraint: binding_id_with_constraints.visibility_constraint,
-            })
-    }
-}
-
-impl std::iter::FusedIterator for BindingWithConstraintsIterator<'_, '_> {}
-
-pub(crate) struct BindingWithConstraints<'map, 'db> {
+pub(crate) struct BindingWithConstraints<'db, I> {
     pub(crate) binding: Option<Definition<'db>>,
-    pub(crate) constraints: ConstraintsIterator<'map, 'db>,
+    pub(crate) constraints: I,
     pub(crate) visibility_constraint: ScopedVisibilityConstraintId,
 }
 
-pub(crate) struct ConstraintsIterator<'map, 'db> {
-    all_constraints: &'map AllConstraints<'db>,
-    constraint_ids: ConstraintIdIterator<'map>,
-}
-
-impl<'db> Iterator for ConstraintsIterator<'_, 'db> {
-    type Item = Constraint<'db>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.constraint_ids
-            .next()
-            .map(|constraint_id| self.all_constraints[constraint_id])
-    }
-}
-
-impl std::iter::FusedIterator for ConstraintsIterator<'_, '_> {}
-
-pub(crate) struct DeclarationsIterator<'map, 'db> {
-    all_definitions: &'map IndexVec<ScopedDefinitionId, Option<Definition<'db>>>,
-    pub(crate) visibility_constraints: &'map VisibilityConstraints<'db>,
-    inner: DeclarationIdIterator<'map>,
-}
-
 pub(crate) struct DeclarationWithConstraint<'db> {
     pub(crate) declaration: Option<Definition<'db>>,
     pub(crate) visibility_constraint: ScopedVisibilityConstraintId,
 }
 
-impl<'db> Iterator for DeclarationsIterator<'_, 'db> {
-    type Item = DeclarationWithConstraint<'db>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next().map(
-            |DeclarationIdWithConstraint {
-                 definition,
-                 visibility_constraint,
-             }| {
-                DeclarationWithConstraint {
-                    declaration: self.all_definitions[definition],
-                    visibility_constraint,
-                }
-            },
-        )
-    }
-}
-
-impl std::iter::FusedIterator for DeclarationsIterator<'_, '_> {}
-
 /// A snapshot of the definitions and constraints state at a particular point in control flow.
 #[derive(Clone, Debug)]
 pub(super) struct FlowSnapshot {

crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs

@@ -43,12 +43,14 @@
 //!
 //! Tracking live declarations is simpler, since constraints are not involved, but otherwise very
 //! similar to tracking live bindings.
-use crate::semantic_index::use_def::VisibilityConstraints;
 
-use super::bitset::{BitSet, BitSetIterator};
+use itertools::{EitherOrBoth, Itertools};
 use ruff_index::newtype_index;
 use smallvec::SmallVec;
 
+use crate::semantic_index::use_def::bitset::BitSet;
+use crate::semantic_index::use_def::VisibilityConstraints;
+
 /// A newtype-index for a definition in a particular scope.
 #[newtype_index]
 pub(super) struct ScopedDefinitionId;
@@ -71,14 +73,12 @@ const INLINE_BINDING_BLOCKS: usize = 3;
 
 /// A [`BitSet`] of [`ScopedDefinitionId`], representing live bindings of a symbol in a scope.
 type Bindings = BitSet<INLINE_BINDING_BLOCKS>;
-type BindingsIterator<'a> = BitSetIterator<'a, INLINE_BINDING_BLOCKS>;
 
 /// Can reference this * 64 total declarations inline; more will fall back to the heap.
 const INLINE_DECLARATION_BLOCKS: usize = 3;
 
 /// A [`BitSet`] of [`ScopedDefinitionId`], representing live declarations of a symbol in a scope.
 type Declarations = BitSet<INLINE_DECLARATION_BLOCKS>;
-type DeclarationsIterator<'a> = BitSetIterator<'a, INLINE_DECLARATION_BLOCKS>;
 
 /// Can reference this * 64 total constraints inline; more will fall back to the heap.
 const INLINE_CONSTRAINT_BLOCKS: usize = 2;
@@ -94,10 +94,6 @@ type InlineConstraintArray = [Constraints; INLINE_BINDINGS_PER_SYMBOL];
 /// One [`BitSet`] of applicable [`ScopedConstraintId`]s per live binding.
 type ConstraintsPerBinding = SmallVec<InlineConstraintArray>;
 
-/// Iterate over all constraints for a single binding.
-type ConstraintsIterator<'a> = std::slice::Iter<'a, Constraints>;
-type ConstraintsIntoIterator = smallvec::IntoIter<InlineConstraintArray>;
-
 /// A newtype-index for a visibility constraint in a particular scope.
 #[newtype_index]
 pub(crate) struct ScopedVisibilityConstraintId;
@@ -120,16 +116,18 @@ type VisibilityConstraintPerDeclaration = SmallVec<InlineVisibilityConstraintsAr
 /// One [`ScopedVisibilityConstraintId`] per live binding.
 type VisibilityConstraintPerBinding = SmallVec<InlineVisibilityConstraintsArray>;
 
-/// Iterator over the visibility constraints for all live bindings/declarations.
-type VisibilityConstraintsIterator<'a> = std::slice::Iter<'a, ScopedVisibilityConstraintId>;
-
-type VisibilityConstraintsIntoIterator = smallvec::IntoIter<InlineVisibilityConstraintsArray>;
-
 /// Live declarations for a single symbol at some point in control flow, with their
 /// corresponding visibility constraints.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub(super) struct SymbolDeclarations {
     /// [`BitSet`]: which declarations (as [`ScopedDefinitionId`]) can reach the current location?
+    ///
+    /// Invariant: Because this is a `BitSet`, it can be viewed as a _sorted_ set of definition
+    /// IDs. The `visibility_constraints` field stores constraints for each definition. Therefore
+    /// those fields must always have the same `len()` as `live_declarations`, and the elements
+    /// must appear in the same order.  Effectively, this means that elements must always be added
+    /// in sorted order, or via a binary search that determines the correct place to insert new
+    /// constraints.
     pub(crate) live_declarations: Declarations,
 
     /// For each live declaration, which visibility constraint applies to it?
@@ -167,19 +165,61 @@ impl SymbolDeclarations {
     }
 
     /// Return an iterator over live declarations for this symbol.
-    pub(super) fn iter(&self) -> DeclarationIdIterator {
-        DeclarationIdIterator {
-            declarations: self.live_declarations.iter(),
-            visibility_constraints: self.visibility_constraints.iter(),
+    pub(super) fn iter(&self) -> impl Iterator<Item = DeclarationIdWithConstraint> + '_ {
+        (self.live_declarations.iter())
+            .zip(self.visibility_constraints.iter())
+            .map(
+                |(declaration, &visibility_constraint)| DeclarationIdWithConstraint {
+                    definition: ScopedDefinitionId::from_u32(declaration),
+                    visibility_constraint,
+                },
+            )
+    }
+
+    fn merge(&mut self, b: Self, visibility_constraints: &mut VisibilityConstraints) {
+        let a = std::mem::take(self);
+        self.live_declarations = a.live_declarations.clone();
+        self.live_declarations.union(&b.live_declarations);
+
+        // Invariant: These zips are well-formed since we maintain an invariant that all of our
+        // fields are sets/vecs with the same length.
+        let a = (a.live_declarations.iter()).zip(a.visibility_constraints);
+        let b = (b.live_declarations.iter()).zip(b.visibility_constraints);
+
+        // Invariant: merge_join_by consumes the two iterators in sorted order, which ensures that
+        // the definition IDs and constraints line up correctly in the merged result. If a
+        // definition is found in both `a` and `b`, we compose the constraints from the two paths
+        // in an appropriate way (intersection for narrowing constraints; ternary OR for visibility
+        // constraints). If a definition is found in only one path, it is used as-is.
+        for zipped in a.merge_join_by(b, |(a_decl, _), (b_decl, _)| a_decl.cmp(b_decl)) {
+            match zipped {
+                EitherOrBoth::Both((_, a_vis_constraint), (_, b_vis_constraint)) => {
+                    let vis_constraint = visibility_constraints
+                        .add_or_constraint(a_vis_constraint, b_vis_constraint);
+                    self.visibility_constraints.push(vis_constraint);
+                }
+
+                EitherOrBoth::Left((_, vis_constraint))
+                | EitherOrBoth::Right((_, vis_constraint)) => {
+                    self.visibility_constraints.push(vis_constraint);
+                }
+            }
         }
     }
 }
 
 /// Live bindings for a single symbol at some point in control flow. Each live binding comes
 /// with a set of narrowing constraints and a visibility constraint.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub(super) struct SymbolBindings {
     /// [`BitSet`]: which bindings (as [`ScopedDefinitionId`]) can reach the current location?
+    ///
+    /// Invariant: Because this is a `BitSet`, it can be viewed as a _sorted_ set of definition
+    /// IDs. The `constraints` and `visibility_constraints` field stores constraints for each
+    /// definition. Therefore those fields must always have the same `len()` as
+    /// `live_bindings`, and the elements must appear in the same order.  Effectively, this means
+    /// that elements must always be added in sorted order, or via a binary search that determines
+    /// the correct place to insert new constraints.
     live_bindings: Bindings,
 
     /// For each live binding, which [`ScopedConstraintId`] apply?
@@ -235,11 +275,75 @@ impl SymbolBindings {
     }
 
     /// Iterate over currently live bindings for this symbol
-    pub(super) fn iter(&self) -> BindingIdWithConstraintsIterator {
-        BindingIdWithConstraintsIterator {
-            definitions: self.live_bindings.iter(),
-            constraints: self.constraints.iter(),
-            visibility_constraints: self.visibility_constraints.iter(),
+    pub(super) fn iter(
+        &self,
+    ) -> impl Iterator<Item = BindingIdWithConstraints<impl Iterator<Item = ScopedConstraintId> + '_>> + '_
+    {
+        let i = (self.live_bindings.iter())
+            .zip(self.constraints.iter())
+            .zip(self.visibility_constraints.iter());
+        i.map(
+            |((def, constraints), visibility_constraint_id)| BindingIdWithConstraints {
+                definition: ScopedDefinitionId::from_u32(def),
+                constraint_ids: constraints.iter().map(ScopedConstraintId::from_u32),
+                visibility_constraint: *visibility_constraint_id,
+            },
+        )
+    }
+
+    fn merge(&mut self, mut b: Self, visibility_constraints: &mut VisibilityConstraints) {
+        let mut a = std::mem::take(self);
+        self.live_bindings = a.live_bindings.clone();
+        self.live_bindings.union(&b.live_bindings);
+
+        // Invariant: These zips are well-formed since we maintain an invariant that all of our
+        // fields are sets/vecs with the same length.
+        //
+        // Performance: We iterate over the `constraints` smallvecs via mut reference, because the
+        // individual elements are `BitSet`s (currently 24 bytes in size), and we don't want to
+        // move them by value multiple times during iteration. By iterating by reference, we only
+        // have to copy single pointers around.  In the loop below, the `std::mem::take` calls
+        // specify precisely where we want to move them into the merged `constraints` smallvec.
+        //
+        // We don't need a similar optimization for `visibility_constraints`, since those elements
+        // are 32-bit IndexVec IDs, and so are already cheap to move/copy.
+        let a = (a.live_bindings.iter())
+            .zip(a.constraints.iter_mut())
+            .zip(a.visibility_constraints);
+        let b = (b.live_bindings.iter())
+            .zip(b.constraints.iter_mut())
+            .zip(b.visibility_constraints);
+
+        // Invariant: merge_join_by consumes the two iterators in sorted order, which ensures that
+        // the definition IDs and constraints line up correctly in the merged result. If a
+        // definition is found in both `a` and `b`, we compose the constraints from the two paths
+        // in an appropriate way (intersection for narrowing constraints; ternary OR for visibility
+        // constraints). If a definition is found in only one path, it is used as-is.
+        for zipped in a.merge_join_by(b, |((a_def, _), _), ((b_def, _), _)| a_def.cmp(b_def)) {
+            match zipped {
+                EitherOrBoth::Both(
+                    ((_, a_constraints), a_vis_constraint),
+                    ((_, b_constraints), b_vis_constraint),
+                ) => {
+                    // If the same definition is visible through both paths, any constraint
+                    // that applies on only one path is irrelevant to the resulting type from
+                    // unioning the two paths, so we intersect the constraints.
+                    let constraints = a_constraints;
+                    constraints.intersect(b_constraints);
+                    self.constraints.push(std::mem::take(constraints));
+
+                    // For visibility constraints, we merge them using a ternary OR operation:
+                    let vis_constraint = visibility_constraints
+                        .add_or_constraint(a_vis_constraint, b_vis_constraint);
+                    self.visibility_constraints.push(vis_constraint);
+                }
+
+                EitherOrBoth::Left(((_, constraints), vis_constraint))
+                | EitherOrBoth::Right(((_, constraints), vis_constraint)) => {
+                    self.constraints.push(std::mem::take(constraints));
+                    self.visibility_constraints.push(vis_constraint);
+                }
+            }
         }
     }
 }
@@ -303,202 +407,9 @@ impl SymbolState {
         b: SymbolState,
         visibility_constraints: &mut VisibilityConstraints,
     ) {
-        let mut a = Self {
-            bindings: SymbolBindings {
-                live_bindings: Bindings::default(),
-                constraints: ConstraintsPerBinding::default(),
-                visibility_constraints: VisibilityConstraintPerBinding::default(),
-            },
-            declarations: SymbolDeclarations {
-                live_declarations: self.declarations.live_declarations.clone(),
-                visibility_constraints: VisibilityConstraintPerDeclaration::default(),
-            },
-        };
-
-        std::mem::swap(&mut a, self);
+        self.bindings.merge(b.bindings, visibility_constraints);
         self.declarations
-            .live_declarations
-            .union(&b.declarations.live_declarations);
-
-        let mut a_defs_iter = a.bindings.live_bindings.iter();
-        let mut b_defs_iter = b.bindings.live_bindings.iter();
-        let mut a_constraints_iter = a.bindings.constraints.into_iter();
-        let mut b_constraints_iter = b.bindings.constraints.into_iter();
-        let mut a_vis_constraints_iter = a.bindings.visibility_constraints.into_iter();
-        let mut b_vis_constraints_iter = b.bindings.visibility_constraints.into_iter();
-
-        let mut opt_a_def: Option<u32> = a_defs_iter.next();
-        let mut opt_b_def: Option<u32> = b_defs_iter.next();
-
-        // Iterate through the definitions from `a` and `b`, always processing the lower definition
-        // ID first, and pushing each definition onto the merged `SymbolState` with its
-        // constraints. If a definition is found in both `a` and `b`, push it with the intersection
-        // of the constraints from the two paths; a constraint that applies from only one possible
-        // path is irrelevant.
-
-        // Helper to push `def`, with constraints in `constraints_iter`, onto `self`.
-        let push = |def,
-                    constraints_iter: &mut ConstraintsIntoIterator,
-                    visibility_constraints_iter: &mut VisibilityConstraintsIntoIterator,
-                    merged: &mut Self| {
-            merged.bindings.live_bindings.insert(def);
-            // SAFETY: we only ever create SymbolState using [`SymbolState::undefined`], which adds
-            // one "unbound" definition with corresponding narrowing and visibility constraints, or
-            // using [`SymbolState::record_binding`] or [`SymbolState::record_declaration`], which
-            // similarly add one definition with corresponding constraints. [`SymbolState::merge`]
-            // always pushes one definition and one constraint bitset and one visibility constraint
-            // together (just below), so the number of definitions and the number of constraints can
-            // never get out of sync.
-            // get out of sync.
-            let constraints = constraints_iter
-                .next()
-                .expect("definitions and constraints length mismatch");
-            let visibility_constraints = visibility_constraints_iter
-                .next()
-                .expect("definitions and visibility_constraints length mismatch");
-            merged.bindings.constraints.push(constraints);
-            merged
-                .bindings
-                .visibility_constraints
-                .push(visibility_constraints);
-        };
-
-        loop {
-            match (opt_a_def, opt_b_def) {
-                (Some(a_def), Some(b_def)) => match a_def.cmp(&b_def) {
-                    std::cmp::Ordering::Less => {
-                        // Next definition ID is only in `a`, push it to `self` and advance `a`.
-                        push(
-                            a_def,
-                            &mut a_constraints_iter,
-                            &mut a_vis_constraints_iter,
-                            self,
-                        );
-                        opt_a_def = a_defs_iter.next();
-                    }
-                    std::cmp::Ordering::Greater => {
-                        // Next definition ID is only in `b`, push it to `self` and advance `b`.
-                        push(
-                            b_def,
-                            &mut b_constraints_iter,
-                            &mut b_vis_constraints_iter,
-                            self,
-                        );
-                        opt_b_def = b_defs_iter.next();
-                    }
-                    std::cmp::Ordering::Equal => {
-                        // Next definition is in both; push to `self` and intersect constraints.
-                        push(
-                            a_def,
-                            &mut b_constraints_iter,
-                            &mut b_vis_constraints_iter,
-                            self,
-                        );
-
-                        // SAFETY: see comment in `push` above.
-                        let a_constraints = a_constraints_iter
-                            .next()
-                            .expect("definitions and constraints length mismatch");
-                        let current_constraints = self.bindings.constraints.last_mut().unwrap();
-
-                        // If the same definition is visible through both paths, any constraint
-                        // that applies on only one path is irrelevant to the resulting type from
-                        // unioning the two paths, so we intersect the constraints.
-                        current_constraints.intersect(&a_constraints);
-
-                        // For visibility constraints, we merge them using a ternary OR operation:
-                        let a_vis_constraint = a_vis_constraints_iter
-                            .next()
-                            .expect("visibility_constraints length mismatch");
-                        let current_vis_constraint =
-                            self.bindings.visibility_constraints.last_mut().unwrap();
-                        *current_vis_constraint = visibility_constraints
-                            .add_or_constraint(*current_vis_constraint, a_vis_constraint);
-
-                        opt_a_def = a_defs_iter.next();
-                        opt_b_def = b_defs_iter.next();
-                    }
-                },
-                (Some(a_def), None) => {
-                    // We've exhausted `b`, just push the def from `a` and move on to the next.
-                    push(
-                        a_def,
-                        &mut a_constraints_iter,
-                        &mut a_vis_constraints_iter,
-                        self,
-                    );
-                    opt_a_def = a_defs_iter.next();
-                }
-                (None, Some(b_def)) => {
-                    // We've exhausted `a`, just push the def from `b` and move on to the next.
-                    push(
-                        b_def,
-                        &mut b_constraints_iter,
-                        &mut b_vis_constraints_iter,
-                        self,
-                    );
-                    opt_b_def = b_defs_iter.next();
-                }
-                (None, None) => break,
-            }
-        }
-
-        // Same as above, but for declarations.
-        let mut a_decls_iter = a.declarations.live_declarations.iter();
-        let mut b_decls_iter = b.declarations.live_declarations.iter();
-        let mut a_vis_constraints_iter = a.declarations.visibility_constraints.into_iter();
-        let mut b_vis_constraints_iter = b.declarations.visibility_constraints.into_iter();
-
-        let mut opt_a_decl: Option<u32> = a_decls_iter.next();
-        let mut opt_b_decl: Option<u32> = b_decls_iter.next();
-
-        let push = |vis_constraints_iter: &mut VisibilityConstraintsIntoIterator,
-                    merged: &mut Self| {
-            let vis_constraints = vis_constraints_iter
-                .next()
-                .expect("declarations and visibility_constraints length mismatch");
-            merged
-                .declarations
-                .visibility_constraints
-                .push(vis_constraints);
-        };
-
-        loop {
-            match (opt_a_decl, opt_b_decl) {
-                (Some(a_decl), Some(b_decl)) => match a_decl.cmp(&b_decl) {
-                    std::cmp::Ordering::Less => {
-                        push(&mut a_vis_constraints_iter, self);
-                        opt_a_decl = a_decls_iter.next();
-                    }
-                    std::cmp::Ordering::Greater => {
-                        push(&mut b_vis_constraints_iter, self);
-                        opt_b_decl = b_decls_iter.next();
-                    }
-                    std::cmp::Ordering::Equal => {
-                        push(&mut b_vis_constraints_iter, self);
-
-                        let a_vis_constraint = a_vis_constraints_iter
-                            .next()
-                            .expect("declarations and visibility_constraints length mismatch");
-                        let current = self.declarations.visibility_constraints.last_mut().unwrap();
-                        *current =
-                            visibility_constraints.add_or_constraint(*current, a_vis_constraint);
-
-                        opt_a_decl = a_decls_iter.next();
-                        opt_b_decl = b_decls_iter.next();
-                    }
-                },
-                (Some(_), None) => {
-                    push(&mut a_vis_constraints_iter, self);
-                    opt_a_decl = a_decls_iter.next();
-                }
-                (None, Some(_)) => {
-                    push(&mut b_vis_constraints_iter, self);
-                    opt_b_decl = b_decls_iter.next();
-                }
-                (None, None) => break,
-            }
-        }
+            .merge(b.declarations, visibility_constraints);
     }
 
     pub(super) fn bindings(&self) -> &SymbolBindings {
@@ -514,61 +425,12 @@ impl SymbolState {
 /// narrowing constraints ([`ScopedConstraintId`]) and a corresponding visibility
 /// visibility constraint ([`ScopedVisibilityConstraintId`]).
 #[derive(Debug)]
-pub(super) struct BindingIdWithConstraints<'map> {
+pub(super) struct BindingIdWithConstraints<I> {
     pub(super) definition: ScopedDefinitionId,
-    pub(super) constraint_ids: ConstraintIdIterator<'map>,
+    pub(super) constraint_ids: I,
     pub(super) visibility_constraint: ScopedVisibilityConstraintId,
 }
 
-#[derive(Debug)]
-pub(super) struct BindingIdWithConstraintsIterator<'map> {
-    definitions: BindingsIterator<'map>,
-    constraints: ConstraintsIterator<'map>,
-    visibility_constraints: VisibilityConstraintsIterator<'map>,
-}
-
-impl<'map> Iterator for BindingIdWithConstraintsIterator<'map> {
-    type Item = BindingIdWithConstraints<'map>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match (
-            self.definitions.next(),
-            self.constraints.next(),
-            self.visibility_constraints.next(),
-        ) {
-            (None, None, None) => None,
-            (Some(def), Some(constraints), Some(visibility_constraint_id)) => {
-                Some(BindingIdWithConstraints {
-                    definition: ScopedDefinitionId::from_u32(def),
-                    constraint_ids: ConstraintIdIterator {
-                        wrapped: constraints.iter(),
-                    },
-                    visibility_constraint: *visibility_constraint_id,
-                })
-            }
-            // SAFETY: see above.
-            _ => unreachable!("definitions and constraints length mismatch"),
-        }
-    }
-}
-
-impl std::iter::FusedIterator for BindingIdWithConstraintsIterator<'_> {}
-
-#[derive(Debug)]
-pub(super) struct ConstraintIdIterator<'a> {
-    wrapped: BitSetIterator<'a, INLINE_CONSTRAINT_BLOCKS>,
-}
-
-impl Iterator for ConstraintIdIterator<'_> {
-    type Item = ScopedConstraintId;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.wrapped.next().map(ScopedConstraintId::from_u32)
-    }
-}
-
-impl std::iter::FusedIterator for ConstraintIdIterator<'_> {}
-
 /// A single declaration (as [`ScopedDefinitionId`]) with a corresponding visibility
 /// visibility constraint ([`ScopedVisibilityConstraintId`]).
 #[derive(Debug)]
@@ -577,31 +439,6 @@ pub(super) struct DeclarationIdWithConstraint {
     pub(super) visibility_constraint: ScopedVisibilityConstraintId,
 }
 
-pub(super) struct DeclarationIdIterator<'map> {
-    pub(crate) declarations: DeclarationsIterator<'map>,
-    pub(crate) visibility_constraints: VisibilityConstraintsIterator<'map>,
-}
-
-impl Iterator for DeclarationIdIterator<'_> {
-    type Item = DeclarationIdWithConstraint;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match (self.declarations.next(), self.visibility_constraints.next()) {
-            (None, None) => None,
-            (Some(declaration), Some(&visibility_constraint)) => {
-                Some(DeclarationIdWithConstraint {
-                    definition: ScopedDefinitionId::from_u32(declaration),
-                    visibility_constraint,
-                })
-            }
-            // SAFETY: see above.
-            _ => unreachable!("declarations and visibility_constraints length mismatch"),
-        }
-    }
-}
-
-impl std::iter::FusedIterator for DeclarationIdIterator<'_> {}
-
 #[cfg(test)]
 mod tests {
     use super::*;

crates/red_knot_python_semantic/src/suppression.rs

@@ -163,6 +163,17 @@ fn check_unknown_rule(context: &mut CheckSuppressionsContext) {
                     format_args!("Unknown rule `{rule}`"),
                 );
             }
+
+            GetLintError::PrefixedWithCategory {
+                prefixed,
+                suggestion,
+            } => {
+                context.report_lint(
+                    &UNKNOWN_RULE,
+                    unknown.range,
+                    format_args!("Unknown rule `{prefixed}`. Did you mean `{suggestion}`?"),
+                );
+            }
         };
     }
 }
@@ -203,7 +214,7 @@ fn check_unused_suppressions(context: &mut CheckSuppressionsContext) {
     );
 
     // Collect all suppressions that are unused after type-checking.
-    for suppression in all {
+    for suppression in all.iter() {
         if context.diagnostics.is_used(suppression.id()) {
             continue;
         }
@@ -390,23 +401,11 @@ impl Suppressions {
             })
     }
 
-    fn iter(&self) -> SuppressionsIter {
+    fn iter(&self) -> impl Iterator<Item = &Suppression> + '_ {
         self.file.iter().chain(&self.line)
     }
 }
 
-pub(crate) type SuppressionsIter<'a> =
-    std::iter::Chain<std::slice::Iter<'a, Suppression>, std::slice::Iter<'a, Suppression>>;
-
-impl<'a> IntoIterator for &'a Suppressions {
-    type Item = &'a Suppression;
-    type IntoIter = SuppressionsIter<'a>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.iter()
-    }
-}
-
 /// A `type: ignore` or `knot: ignore` suppression.
 ///
 /// Suppression comments that suppress multiple codes
@@ -765,8 +764,9 @@ impl<'src> SuppressionParser<'src> {
 
     fn eat_word(&mut self) -> bool {
         if self.cursor.eat_if(char::is_alphabetic) {
+            // Allow `:` for better error recovery when someone uses `lint:code` instead of just `code`.
             self.cursor
-                .eat_while(|c| c.is_alphanumeric() || matches!(c, '_' | '-'));
+                .eat_while(|c| c.is_alphanumeric() || matches!(c, '_' | '-' | ':'));
             true
         } else {
             false

crates/red_knot_python_semantic/src/types.rs

@@ -23,12 +23,12 @@ pub use self::subclass_of::SubclassOfType;
 use crate::module_name::ModuleName;
 use crate::module_resolver::{file_to_module, resolve_module, KnownModule};
 use crate::semantic_index::ast_ids::HasScopedExpressionId;
+use crate::semantic_index::constraint::Constraint;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId};
 use crate::semantic_index::{
     global_scope, imported_modules, semantic_index, symbol_table, use_def_map,
-    BindingWithConstraints, BindingWithConstraintsIterator, DeclarationWithConstraint,
-    DeclarationsIterator,
+    BindingWithConstraints, DeclarationWithConstraint, UseDefMap,
 };
 use crate::stdlib::{builtins_symbol, known_module_symbol, typing_extensions_symbol};
 use crate::suppression::check_suppressions;
@@ -118,7 +118,7 @@ fn symbol<'db>(db: &'db dyn Db, scope: ScopeId<'db>, name: &str) -> Symbol<'db>
         // on inference from bindings.
 
         let declarations = use_def.public_declarations(symbol_id);
-        let declared = symbol_from_declarations(db, declarations);
+        let declared = symbol_from_declarations(db, use_def.as_ref(), declarations);
         let is_final = declared.as_ref().is_ok_and(SymbolAndQualifiers::is_final);
         let declared = declared.map(|SymbolAndQualifiers(symbol, _)| symbol);
 
@@ -128,7 +128,7 @@ fn symbol<'db>(db: &'db dyn Db, scope: ScopeId<'db>, name: &str) -> Symbol<'db>
             // Symbol is possibly declared
             Ok(Symbol::Type(declared_ty, Boundness::PossiblyUnbound)) => {
                 let bindings = use_def.public_bindings(symbol_id);
-                let inferred = symbol_from_bindings(db, bindings);
+                let inferred = symbol_from_bindings(db, use_def.as_ref(), bindings);
 
                 match inferred {
                     // Symbol is possibly undeclared and definitely unbound
@@ -148,7 +148,7 @@ fn symbol<'db>(db: &'db dyn Db, scope: ScopeId<'db>, name: &str) -> Symbol<'db>
             // Symbol is undeclared, return the union of `Unknown` with the inferred type
             Ok(Symbol::Unbound) => {
                 let bindings = use_def.public_bindings(symbol_id);
-                let inferred = symbol_from_bindings(db, bindings);
+                let inferred = symbol_from_bindings(db, use_def.as_ref(), bindings);
 
                 widen_type_for_undeclared_public_symbol(db, inferred, is_dunder_slots || is_final)
             }
@@ -321,11 +321,14 @@ fn definition_expression_type<'db>(
 /// together with boundness information in a [`Symbol`].
 ///
 /// The type will be a union if there are multiple bindings with different types.
-fn symbol_from_bindings<'db>(
+fn symbol_from_bindings<'map, 'db: 'map>(
     db: &'db dyn Db,
-    bindings_with_constraints: BindingWithConstraintsIterator<'_, 'db>,
+    use_def: &UseDefMap<'map>,
+    bindings_with_constraints: impl Iterator<
+        Item = BindingWithConstraints<'db, impl Iterator<Item = Constraint<'db>>>,
+    >,
 ) -> Symbol<'db> {
-    let visibility_constraints = bindings_with_constraints.visibility_constraints;
+    let visibility_constraints = &use_def.visibility_constraints;
     let mut bindings_with_constraints = bindings_with_constraints.peekable();
 
     let unbound_visibility = if let Some(BindingWithConstraints {
@@ -433,7 +436,7 @@ impl<'db> From<Type<'db>> for SymbolAndQualifiers<'db> {
 type SymbolFromDeclarationsResult<'db> =
     Result<SymbolAndQualifiers<'db>, (TypeAndQualifiers<'db>, Box<[Type<'db>]>)>;
 
-/// Build a declared type from a [`DeclarationsIterator`].
+/// Build a declared type from an iterator of [`DeclarationWithConstraint`]s.
 ///
 /// If there is only one declaration, or all declarations declare the same type, returns
 /// `Ok(..)`. If there are conflicting declarations, returns an `Err(..)` variant with
@@ -443,9 +446,10 @@ type SymbolFromDeclarationsResult<'db> =
 /// [`TypeQualifiers`] that have been specified on the declaration(s).
 fn symbol_from_declarations<'db>(
     db: &'db dyn Db,
-    declarations: DeclarationsIterator<'_, 'db>,
+    use_def: &UseDefMap,
+    declarations: impl Iterator<Item = DeclarationWithConstraint<'db>>,
 ) -> SymbolFromDeclarationsResult<'db> {
-    let visibility_constraints = declarations.visibility_constraints;
+    let visibility_constraints = &use_def.visibility_constraints;
     let mut declarations = declarations.peekable();
 
     let undeclared_visibility = if let Some(DeclarationWithConstraint {
@@ -4126,7 +4130,7 @@ impl<'db> Class<'db> {
 
             let declarations = use_def.public_declarations(symbol_id);
 
-            match symbol_from_declarations(db, declarations) {
+            match symbol_from_declarations(db, use_def.as_ref(), declarations) {
                 Ok(SymbolAndQualifiers(Symbol::Type(declared_ty, _), qualifiers)) => {
                     if let Some(function) = declared_ty.into_function_literal() {
                         // TODO: Eventually, we are going to process all decorators correctly. This is
@@ -4143,7 +4147,7 @@ impl<'db> Class<'db> {
                 }
                 Ok(symbol @ SymbolAndQualifiers(Symbol::Unbound, qualifiers)) => {
                     let bindings = use_def.public_bindings(symbol_id);
-                    let inferred = symbol_from_bindings(db, bindings);
+                    let inferred = symbol_from_bindings(db, use_def.as_ref(), bindings);
 
                     SymbolAndQualifiers(
                         widen_type_for_undeclared_public_symbol(db, inferred, symbol.is_final()),

crates/red_knot_python_semantic/src/types/infer.rs

@@ -862,7 +862,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         let use_def = self.index.use_def_map(binding.file_scope(self.db()));
         let declarations = use_def.declarations_at_binding(binding);
         let mut bound_ty = ty;
-        let declared_ty = symbol_from_declarations(self.db(), declarations)
+        let declared_ty = symbol_from_declarations(self.db(), use_def.as_ref(), declarations)
             .map(|SymbolAndQualifiers(s, _)| s.ignore_possibly_unbound().unwrap_or(Type::unknown()))
             .unwrap_or_else(|(ty, conflicting)| {
                 // TODO point out the conflicting declarations in the diagnostic?
@@ -897,7 +897,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         let use_def = self.index.use_def_map(declaration.file_scope(self.db()));
         let prior_bindings = use_def.bindings_at_declaration(declaration);
         // unbound_ty is Never because for this check we don't care about unbound
-        let inferred_ty = symbol_from_bindings(self.db(), prior_bindings)
+        let inferred_ty = symbol_from_bindings(self.db(), use_def.as_ref(), prior_bindings)
             .ignore_possibly_unbound()
             .unwrap_or(Type::Never);
         let ty = if inferred_ty.is_assignable_to(self.db(), ty.inner_type()) {
@@ -3362,7 +3362,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         // If we're inferring types of deferred expressions, always treat them as public symbols
         let inferred = if self.is_deferred() {
             if let Some(symbol) = self.index.symbol_table(file_scope_id).symbol_id_by_name(id) {
-                symbol_from_bindings(self.db(), use_def.public_bindings(symbol))
+                symbol_from_bindings(self.db(), use_def.as_ref(), use_def.public_bindings(symbol))
             } else {
                 assert!(
                     self.deferred_state.in_string_annotation(),
@@ -3372,7 +3372,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             }
         } else {
             let use_id = name.scoped_use_id(self.db(), self.scope());
-            symbol_from_bindings(self.db(), use_def.bindings_at_use(use_id))
+            symbol_from_bindings(self.db(), use_def.as_ref(), use_def.bindings_at_use(use_id))
         };
 
         if let Symbol::Type(ty, Boundness::Bound) = inferred {

crates/ruff_db/src/diagnostic.rs

@@ -94,6 +94,10 @@ impl DiagnosticId {
         matches!(self, DiagnosticId::Lint(self_name) if self_name == name)
     }
 
+    pub fn strip_category(code: &str) -> Option<&str> {
+        code.split_once(':').map(|(_, rest)| rest)
+    }
+
     /// Returns `true` if this `DiagnosticId` matches the given name.
     ///
     /// ## Examples

crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py

@@ -1,8 +1,7 @@
-from __future__ import annotations
-
 from datetime import datetime
 
 import pendulum
+
 from airflow.decorators import dag, task
 from airflow.models import DAG
 from airflow.models.baseoperator import BaseOperator
@@ -14,22 +13,30 @@ from airflow.utils.context import get_current_context
 
 def access_invalid_key_in_context(**context):
     print("access invalid key", context["conf"])
-    print("access invalid key", context.get("conf"))
-
 
 @task
 def access_invalid_key_task_out_of_dag(**context):
-    print("access invalid key", context["conf"])
     print("access invalid key", context.get("conf"))
 
+@dag(
+    schedule=None,
+    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    catchup=False,
+    tags=[""],
+)
+def invalid_dag():
+    @task()
+    def access_invalid_key_task(**context):
+        print("access invalid key", context.get("conf"))
 
-@task
-def access_invalid_argument_task_out_of_dag(
-    execution_date, tomorrow_ds, logical_date, **context
-):
-    print("execution date", execution_date)
-    print("access invalid key", context.get("conf"))
+    task1 = PythonOperator(
+        task_id="task1",
+        python_callable=access_invalid_key_in_context,
+    )
+    access_invalid_key_task() >> task1
+    access_invalid_key_task_out_of_dag()
 
+invalid_dag()
 
 @task
 def print_config(**context):
@@ -49,9 +56,31 @@ def print_config(**context):
     yesterday_ds = context["yesterday_ds"]
     yesterday_ds_nodash = context["yesterday_ds_nodash"]
 
+with DAG(
+    dag_id="example_dag",
+    schedule_interval="@daily",
+    start_date=datetime(2023, 1, 1),
+    template_searchpath=["/templates"],
+) as dag:
+    task1 = DummyOperator(
+        task_id="task1",
+        params={
+            # Removed variables in template
+            "execution_date": "{{ execution_date }}",
+            "next_ds": "{{ next_ds }}",
+            "prev_ds": "{{ prev_ds }}"
+        },
+    )
+
+class CustomMacrosPlugin(AirflowPlugin):
+    name = "custom_macros"
+    macros = {
+        "execution_date_macro": lambda context: context["execution_date"],
+        "next_ds_macro": lambda context: context["next_ds"]
+    }
 
 @task
-def print_config_with_get_current_context():
+def print_config():
     context = get_current_context()
     execution_date = context["execution_date"]
     next_ds = context["next_ds"]
@@ -65,74 +94,8 @@ def print_config_with_get_current_context():
     yesterday_ds = context["yesterday_ds"]
     yesterday_ds_nodash = context["yesterday_ds_nodash"]
 
-
-@task(task_id="print_the_context")
-def print_context(ds=None, **kwargs):
-    """Print the Airflow context and ds variable from the context."""
-    print(ds)
-    print(kwargs.get("tomorrow_ds"))
-    c = get_current_context()
-    c.get("execution_date")
-
-
-@dag(
-    schedule=None,
-    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
-    catchup=False,
-    tags=[""],
-)
-def invalid_dag():
-    @task()
-    def access_invalid_key_task(**context):
-        print("access invalid key", context.get("conf"))
-
-    @task()
-    def access_invalid_key_explicit_task(execution_date):
-        print(execution_date)
-
-    task1 = PythonOperator(
-        task_id="task1",
-        python_callable=access_invalid_key_in_context,
-    )
-
-    access_invalid_key_task() >> task1
-    access_invalid_key_explicit_task()
-    access_invalid_argument_task_out_of_dag()
-    access_invalid_key_task_out_of_dag()
-    print_config()
-    print_config_with_get_current_context()
-    print_context()
-
-
-invalid_dag()
-
-with DAG(
-    dag_id="example_dag",
-    schedule_interval="@daily",
-    start_date=datetime(2023, 1, 1),
-    template_searchpath=["/templates"],
-) as dag:
-    task1 = DummyOperator(
-        task_id="task1",
-        params={
-            # Removed variables in template
-            "execution_date": "{{ execution_date }}",
-            "next_ds": "{{ next_ds }}",
-            "prev_ds": "{{ prev_ds }}",
-        },
-    )
-
-
-class CustomMacrosPlugin(AirflowPlugin):
-    name = "custom_macros"
-    macros = {
-        "execution_date_macro": lambda context: context["execution_date"],
-        "next_ds_macro": lambda context: context["next_ds"],
-    }
-
-
 class CustomOperator(BaseOperator):
-    def execute(self, next_ds, context):
+    def execute(self, context):
         execution_date = context["execution_date"]
         next_ds = context["next_ds"]
         next_ds_nodash = context["next_ds_nodash"]
@@ -145,6 +108,18 @@ class CustomOperator(BaseOperator):
         yesterday_ds = context["yesterday_ds"]
         yesterday_ds_nodash = context["yesterday_ds_nodash"]
 
+@task
+def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
+    print("execution date", execution_date)
+    print("access invalid key", context.get("conf"))
+
+@task(task_id="print_the_context")
+def print_context(ds=None, **kwargs):
+    """Print the Airflow context and ds variable from the context."""
+    print(ds)
+    print(kwargs.get("tomorrow_ds"))
+    c = get_current_context()
+    c.get("execution_date")
 
 class CustomOperatorNew(BaseOperator):
     def execute(self, context):

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -176,7 +176,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 pyupgrade::rules::use_pep646_unpack(checker, subscript);
             }
             if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::airflow_3_removal_expr(checker, expr);
+                airflow::rules::removed_in_3(checker, expr);
             }
             pandas_vet::rules::subscript(checker, value, expr);
         }
@@ -227,7 +227,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         refurb::rules::regex_flag_alias(checker, expr);
                     }
                     if checker.enabled(Rule::Airflow3Removal) {
-                        airflow::rules::airflow_3_removal_expr(checker, expr);
+                        airflow::rules::removed_in_3(checker, expr);
                     }
                     if checker.enabled(Rule::Airflow3MovedToProvider) {
                         airflow::rules::moved_to_provider_in_3(checker, expr);
@@ -311,7 +311,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         }
                     }
                     if checker.enabled(Rule::Airflow3Removal) {
-                        airflow::rules::airflow_3_removal_expr(checker, expr);
+                        airflow::rules::removed_in_3(checker, expr);
                     }
                     if checker.enabled(Rule::MixedCaseVariableInGlobalScope) {
                         if matches!(checker.semantic.current_scope().kind, ScopeKind::Module) {
@@ -449,7 +449,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 flake8_pyi::rules::bytestring_attribute(checker, expr);
             }
             if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::airflow_3_removal_expr(checker, expr);
+                airflow::rules::removed_in_3(checker, expr);
             }
         }
         Expr::Call(
@@ -1150,7 +1150,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 ruff::rules::unnecessary_regular_expression(checker, call);
             }
             if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::airflow_3_removal_expr(checker, expr);
+                airflow::rules::removed_in_3(checker, expr);
             }
             if checker.enabled(Rule::UnnecessaryCastToInt) {
                 ruff::rules::unnecessary_cast_to_int(checker, call);

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -377,7 +377,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 flake8_pytest_style::rules::parameter_with_default_argument(checker, function_def);
             }
             if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::airflow_3_removal_function_def(checker, function_def);
+                airflow::rules::removed_in_3_function_def(checker, function_def);
             }
             if checker.enabled(Rule::NonPEP695GenericFunction) {
                 pyupgrade::rules::non_pep695_generic_function(checker, function_def);

crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs

@@ -80,7 +80,7 @@ enum Replacement {
 }
 
 /// AIR302
-pub(crate) fn airflow_3_removal_expr(checker: &mut Checker, expr: &Expr) {
+pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
         return;
     }
@@ -117,10 +117,7 @@ pub(crate) fn airflow_3_removal_expr(checker: &mut Checker, expr: &Expr) {
 }
 
 /// AIR302
-pub(crate) fn airflow_3_removal_function_def(
-    checker: &mut Checker,
-    function_def: &StmtFunctionDef,
-) {
+pub(crate) fn removed_in_3_function_def(checker: &mut Checker, function_def: &StmtFunctionDef) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
         return;
     }
@@ -157,9 +154,7 @@ const REMOVED_CONTEXT_KEYS: [&str; 12] = [
 ///     pass
 /// ```
 fn check_function_parameters(checker: &mut Checker, function_def: &StmtFunctionDef) {
-    if !is_airflow_task(function_def, checker.semantic())
-        && !is_execute_method_inherits_from_airflow_operator(function_def, checker.semantic())
-    {
+    if !is_airflow_task(function_def, checker.semantic()) {
         return;
     }
 
@@ -1081,35 +1076,3 @@ fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) ->
             })
     })
 }
-
-/// Check it's "execute" method inherits from Airflow base operator
-///
-/// For example:
-///
-/// ```python
-/// from airflow.models.baseoperator import BaseOperator
-///
-/// class CustomOperator(BaseOperator):
-///     def execute(self):
-///         pass
-/// ```
-fn is_execute_method_inherits_from_airflow_operator(
-    function_def: &StmtFunctionDef,
-    semantic: &SemanticModel,
-) -> bool {
-    if function_def.name.as_str() != "execute" {
-        return false;
-    }
-
-    let ScopeKind::Class(class_def) = semantic.current_scope().kind else {
-        return false;
-    };
-
-    class_def.bases().iter().any(|class_base| {
-        semantic
-            .resolve_qualified_name(class_base)
-            .is_some_and(|qualified_name| {
-                matches!(qualified_name.segments(), ["airflow", .., "BaseOperator"])
-            })
-    })
-}

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap

@@ -1,338 +1,319 @@
 ---
 source: crates/ruff_linter/src/rules/airflow/mod.rs
-snapshot_kind: text
 ---
-AIR302_context.py:22:41: AIR302 `conf` is removed in Airflow 3.0
+AIR302_context.py:19:45: AIR302 `conf` is removed in Airflow 3.0
    |
-20 | @task
-21 | def access_invalid_key_task_out_of_dag(**context):
-22 |     print("access invalid key", context["conf"])
-   |                                         ^^^^^^ AIR302
-23 |     print("access invalid key", context.get("conf"))
-   |
-
-AIR302_context.py:23:45: AIR302 `conf` is removed in Airflow 3.0
-   |
-21 | def access_invalid_key_task_out_of_dag(**context):
-22 |     print("access invalid key", context["conf"])
-23 |     print("access invalid key", context.get("conf"))
+17 | @task
+18 | def access_invalid_key_task_out_of_dag(**context):
+19 |     print("access invalid key", context.get("conf"))
    |                                             ^^^^^^ AIR302
+20 |
+21 | @dag(
    |
 
-AIR302_context.py:28:5: AIR302 `execution_date` is removed in Airflow 3.0
+AIR302_context.py:30:49: AIR302 `conf` is removed in Airflow 3.0
    |
-26 | @task
-27 | def access_invalid_argument_task_out_of_dag(
-28 |     execution_date, tomorrow_ds, logical_date, **context
-   |     ^^^^^^^^^^^^^^ AIR302
-29 | ):
-30 |     print("execution date", execution_date)
-   |
-
-AIR302_context.py:28:21: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-26 | @task
-27 | def access_invalid_argument_task_out_of_dag(
-28 |     execution_date, tomorrow_ds, logical_date, **context
-   |                     ^^^^^^^^^^^ AIR302
-29 | ):
-30 |     print("execution date", execution_date)
-   |
-
-AIR302_context.py:31:45: AIR302 `conf` is removed in Airflow 3.0
-   |
-29 | ):
-30 |     print("execution date", execution_date)
-31 |     print("access invalid key", context.get("conf"))
-   |                                             ^^^^^^ AIR302
-   |
-
-AIR302_context.py:40:30: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-39 |     # Removed usage - should trigger violations
-40 |     execution_date = context["execution_date"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-41 |     next_ds = context["next_ds"]
-42 |     next_ds_nodash = context["next_ds_nodash"]
-   |
-
-AIR302_context.py:41:23: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-39 |     # Removed usage - should trigger violations
-40 |     execution_date = context["execution_date"]
-41 |     next_ds = context["next_ds"]
-   |                       ^^^^^^^^^ AIR302
-42 |     next_ds_nodash = context["next_ds_nodash"]
-43 |     next_execution_date = context["next_execution_date"]
-   |
-
-AIR302_context.py:42:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
-   |
-40 |     execution_date = context["execution_date"]
-41 |     next_ds = context["next_ds"]
-42 |     next_ds_nodash = context["next_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-43 |     next_execution_date = context["next_execution_date"]
-44 |     prev_ds = context["prev_ds"]
-   |
-
-AIR302_context.py:43:35: AIR302 `next_execution_date` is removed in Airflow 3.0
-   |
-41 |     next_ds = context["next_ds"]
-42 |     next_ds_nodash = context["next_ds_nodash"]
-43 |     next_execution_date = context["next_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-44 |     prev_ds = context["prev_ds"]
-45 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |
-
-AIR302_context.py:44:23: AIR302 `prev_ds` is removed in Airflow 3.0
-   |
-42 |     next_ds_nodash = context["next_ds_nodash"]
-43 |     next_execution_date = context["next_execution_date"]
-44 |     prev_ds = context["prev_ds"]
-   |                       ^^^^^^^^^ AIR302
-45 |     prev_ds_nodash = context["prev_ds_nodash"]
-46 |     prev_execution_date = context["prev_execution_date"]
-   |
-
-AIR302_context.py:45:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
-   |
-43 |     next_execution_date = context["next_execution_date"]
-44 |     prev_ds = context["prev_ds"]
-45 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-46 |     prev_execution_date = context["prev_execution_date"]
-47 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |
-
-AIR302_context.py:46:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
-   |
-44 |     prev_ds = context["prev_ds"]
-45 |     prev_ds_nodash = context["prev_ds_nodash"]
-46 |     prev_execution_date = context["prev_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-47 |     prev_execution_date_success = context["prev_execution_date_success"]
-48 |     tomorrow_ds = context["tomorrow_ds"]
-   |
-
-AIR302_context.py:47:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
-   |
-45 |     prev_ds_nodash = context["prev_ds_nodash"]
-46 |     prev_execution_date = context["prev_execution_date"]
-47 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
-48 |     tomorrow_ds = context["tomorrow_ds"]
-49 |     yesterday_ds = context["yesterday_ds"]
-   |
-
-AIR302_context.py:48:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-46 |     prev_execution_date = context["prev_execution_date"]
-47 |     prev_execution_date_success = context["prev_execution_date_success"]
-48 |     tomorrow_ds = context["tomorrow_ds"]
-   |                           ^^^^^^^^^^^^^ AIR302
-49 |     yesterday_ds = context["yesterday_ds"]
-50 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:49:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
-   |
-47 |     prev_execution_date_success = context["prev_execution_date_success"]
-48 |     tomorrow_ds = context["tomorrow_ds"]
-49 |     yesterday_ds = context["yesterday_ds"]
-   |                            ^^^^^^^^^^^^^^ AIR302
-50 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:50:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
-   |
-48 |     tomorrow_ds = context["tomorrow_ds"]
-49 |     yesterday_ds = context["yesterday_ds"]
-50 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-   |
-
-AIR302_context.py:56:30: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-54 | def print_config_with_get_current_context():
-55 |     context = get_current_context()
-56 |     execution_date = context["execution_date"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-57 |     next_ds = context["next_ds"]
-58 |     next_ds_nodash = context["next_ds_nodash"]
-   |
-
-AIR302_context.py:57:23: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-55 |     context = get_current_context()
-56 |     execution_date = context["execution_date"]
-57 |     next_ds = context["next_ds"]
-   |                       ^^^^^^^^^ AIR302
-58 |     next_ds_nodash = context["next_ds_nodash"]
-59 |     next_execution_date = context["next_execution_date"]
-   |
-
-AIR302_context.py:58:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
-   |
-56 |     execution_date = context["execution_date"]
-57 |     next_ds = context["next_ds"]
-58 |     next_ds_nodash = context["next_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-59 |     next_execution_date = context["next_execution_date"]
-60 |     prev_ds = context["prev_ds"]
-   |
-
-AIR302_context.py:59:35: AIR302 `next_execution_date` is removed in Airflow 3.0
-   |
-57 |     next_ds = context["next_ds"]
-58 |     next_ds_nodash = context["next_ds_nodash"]
-59 |     next_execution_date = context["next_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-60 |     prev_ds = context["prev_ds"]
-61 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |
-
-AIR302_context.py:60:23: AIR302 `prev_ds` is removed in Airflow 3.0
-   |
-58 |     next_ds_nodash = context["next_ds_nodash"]
-59 |     next_execution_date = context["next_execution_date"]
-60 |     prev_ds = context["prev_ds"]
-   |                       ^^^^^^^^^ AIR302
-61 |     prev_ds_nodash = context["prev_ds_nodash"]
-62 |     prev_execution_date = context["prev_execution_date"]
-   |
-
-AIR302_context.py:61:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
-   |
-59 |     next_execution_date = context["next_execution_date"]
-60 |     prev_ds = context["prev_ds"]
-61 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-62 |     prev_execution_date = context["prev_execution_date"]
-63 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |
-
-AIR302_context.py:62:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
-   |
-60 |     prev_ds = context["prev_ds"]
-61 |     prev_ds_nodash = context["prev_ds_nodash"]
-62 |     prev_execution_date = context["prev_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-63 |     prev_execution_date_success = context["prev_execution_date_success"]
-64 |     tomorrow_ds = context["tomorrow_ds"]
-   |
-
-AIR302_context.py:63:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
-   |
-61 |     prev_ds_nodash = context["prev_ds_nodash"]
-62 |     prev_execution_date = context["prev_execution_date"]
-63 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
-64 |     tomorrow_ds = context["tomorrow_ds"]
-65 |     yesterday_ds = context["yesterday_ds"]
-   |
-
-AIR302_context.py:64:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-62 |     prev_execution_date = context["prev_execution_date"]
-63 |     prev_execution_date_success = context["prev_execution_date_success"]
-64 |     tomorrow_ds = context["tomorrow_ds"]
-   |                           ^^^^^^^^^^^^^ AIR302
-65 |     yesterday_ds = context["yesterday_ds"]
-66 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:65:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
-   |
-63 |     prev_execution_date_success = context["prev_execution_date_success"]
-64 |     tomorrow_ds = context["tomorrow_ds"]
-65 |     yesterday_ds = context["yesterday_ds"]
-   |                            ^^^^^^^^^^^^^^ AIR302
-66 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:66:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
-   |
-64 |     tomorrow_ds = context["tomorrow_ds"]
-65 |     yesterday_ds = context["yesterday_ds"]
-66 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-   |
-
-AIR302_context.py:73:22: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-71 |     """Print the Airflow context and ds variable from the context."""
-72 |     print(ds)
-73 |     print(kwargs.get("tomorrow_ds"))
-   |                      ^^^^^^^^^^^^^ AIR302
-74 |     c = get_current_context()
-75 |     c.get("execution_date")
-   |
-
-AIR302_context.py:75:11: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-73 |     print(kwargs.get("tomorrow_ds"))
-74 |     c = get_current_context()
-75 |     c.get("execution_date")
-   |           ^^^^^^^^^^^^^^^^ AIR302
-   |
-
-AIR302_context.py:87:49: AIR302 `conf` is removed in Airflow 3.0
-   |
-85 |     @task()
-86 |     def access_invalid_key_task(**context):
-87 |         print("access invalid key", context.get("conf"))
+28 |     @task()
+29 |     def access_invalid_key_task(**context):
+30 |         print("access invalid key", context.get("conf"))
    |                                                 ^^^^^^ AIR302
-88 |
-89 |     @task()
+31 |
+32 |     task1 = PythonOperator(
    |
 
-AIR302_context.py:90:42: AIR302 `execution_date` is removed in Airflow 3.0
+AIR302_context.py:47:30: AIR302 `execution_date` is removed in Airflow 3.0
    |
-89 |     @task()
-90 |     def access_invalid_key_explicit_task(execution_date):
-   |                                          ^^^^^^^^^^^^^^ AIR302
-91 |         print(execution_date)
+46 |     # Removed usage - should trigger violations
+47 |     execution_date = context["execution_date"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+48 |     next_ds = context["next_ds"]
+49 |     next_ds_nodash = context["next_ds_nodash"]
    |
 
-AIR302_context.py:111:5: AIR302 [*] `schedule_interval` is removed in Airflow 3.0
-    |
-109 | with DAG(
-110 |     dag_id="example_dag",
-111 |     schedule_interval="@daily",
-    |     ^^^^^^^^^^^^^^^^^ AIR302
-112 |     start_date=datetime(2023, 1, 1),
-113 |     template_searchpath=["/templates"],
-    |
-    = help: Use `schedule` instead
+AIR302_context.py:48:23: AIR302 `next_ds` is removed in Airflow 3.0
+   |
+46 |     # Removed usage - should trigger violations
+47 |     execution_date = context["execution_date"]
+48 |     next_ds = context["next_ds"]
+   |                       ^^^^^^^^^ AIR302
+49 |     next_ds_nodash = context["next_ds_nodash"]
+50 |     next_execution_date = context["next_execution_date"]
+   |
+
+AIR302_context.py:49:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
+   |
+47 |     execution_date = context["execution_date"]
+48 |     next_ds = context["next_ds"]
+49 |     next_ds_nodash = context["next_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+50 |     next_execution_date = context["next_execution_date"]
+51 |     prev_ds = context["prev_ds"]
+   |
+
+AIR302_context.py:50:35: AIR302 `next_execution_date` is removed in Airflow 3.0
+   |
+48 |     next_ds = context["next_ds"]
+49 |     next_ds_nodash = context["next_ds_nodash"]
+50 |     next_execution_date = context["next_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+51 |     prev_ds = context["prev_ds"]
+52 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |
+
+AIR302_context.py:51:23: AIR302 `prev_ds` is removed in Airflow 3.0
+   |
+49 |     next_ds_nodash = context["next_ds_nodash"]
+50 |     next_execution_date = context["next_execution_date"]
+51 |     prev_ds = context["prev_ds"]
+   |                       ^^^^^^^^^ AIR302
+52 |     prev_ds_nodash = context["prev_ds_nodash"]
+53 |     prev_execution_date = context["prev_execution_date"]
+   |
+
+AIR302_context.py:52:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
+   |
+50 |     next_execution_date = context["next_execution_date"]
+51 |     prev_ds = context["prev_ds"]
+52 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+53 |     prev_execution_date = context["prev_execution_date"]
+54 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |
+
+AIR302_context.py:53:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
+   |
+51 |     prev_ds = context["prev_ds"]
+52 |     prev_ds_nodash = context["prev_ds_nodash"]
+53 |     prev_execution_date = context["prev_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+54 |     prev_execution_date_success = context["prev_execution_date_success"]
+55 |     tomorrow_ds = context["tomorrow_ds"]
+   |
+
+AIR302_context.py:54:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
+   |
+52 |     prev_ds_nodash = context["prev_ds_nodash"]
+53 |     prev_execution_date = context["prev_execution_date"]
+54 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
+55 |     tomorrow_ds = context["tomorrow_ds"]
+56 |     yesterday_ds = context["yesterday_ds"]
+   |
+
+AIR302_context.py:55:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
+   |
+53 |     prev_execution_date = context["prev_execution_date"]
+54 |     prev_execution_date_success = context["prev_execution_date_success"]
+55 |     tomorrow_ds = context["tomorrow_ds"]
+   |                           ^^^^^^^^^^^^^ AIR302
+56 |     yesterday_ds = context["yesterday_ds"]
+57 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:56:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
+   |
+54 |     prev_execution_date_success = context["prev_execution_date_success"]
+55 |     tomorrow_ds = context["tomorrow_ds"]
+56 |     yesterday_ds = context["yesterday_ds"]
+   |                            ^^^^^^^^^^^^^^ AIR302
+57 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:57:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
+   |
+55 |     tomorrow_ds = context["tomorrow_ds"]
+56 |     yesterday_ds = context["yesterday_ds"]
+57 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+58 |
+59 | with DAG(
+   |
+
+AIR302_context.py:61:5: AIR302 [*] `schedule_interval` is removed in Airflow 3.0
+   |
+59 | with DAG(
+60 |     dag_id="example_dag",
+61 |     schedule_interval="@daily",
+   |     ^^^^^^^^^^^^^^^^^ AIR302
+62 |     start_date=datetime(2023, 1, 1),
+63 |     template_searchpath=["/templates"],
+   |
+   = help: Use `schedule` instead
 
 ℹ Safe fix
-108 108 | 
-109 109 | with DAG(
-110 110 |     dag_id="example_dag",
-111     |-    schedule_interval="@daily",
-    111 |+    schedule="@daily",
-112 112 |     start_date=datetime(2023, 1, 1),
-113 113 |     template_searchpath=["/templates"],
-114 114 | ) as dag:
+58 58 | 
+59 59 | with DAG(
+60 60 |     dag_id="example_dag",
+61    |-    schedule_interval="@daily",
+   61 |+    schedule="@daily",
+62 62 |     start_date=datetime(2023, 1, 1),
+63 63 |     template_searchpath=["/templates"],
+64 64 | ) as dag:
 
-AIR302_context.py:115:13: AIR302 `airflow.operators.dummy.DummyOperator` is removed in Airflow 3.0
-    |
-113 |     template_searchpath=["/templates"],
-114 | ) as dag:
-115 |     task1 = DummyOperator(
-    |             ^^^^^^^^^^^^^ AIR302
-116 |         task_id="task1",
-117 |         params={
-    |
-    = help: Use `airflow.operators.empty.EmptyOperator` instead
+AIR302_context.py:65:13: AIR302 `airflow.operators.dummy.DummyOperator` is removed in Airflow 3.0
+   |
+63 |     template_searchpath=["/templates"],
+64 | ) as dag:
+65 |     task1 = DummyOperator(
+   |             ^^^^^^^^^^^^^ AIR302
+66 |         task_id="task1",
+67 |         params={
+   |
+   = help: Use `airflow.operators.empty.EmptyOperator` instead
 
-AIR302_context.py:135:23: AIR302 `next_ds` is removed in Airflow 3.0
+AIR302_context.py:85:30: AIR302 `execution_date` is removed in Airflow 3.0
+   |
+83 | def print_config():
+84 |     context = get_current_context()
+85 |     execution_date = context["execution_date"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+86 |     next_ds = context["next_ds"]
+87 |     next_ds_nodash = context["next_ds_nodash"]
+   |
+
+AIR302_context.py:86:23: AIR302 `next_ds` is removed in Airflow 3.0
+   |
+84 |     context = get_current_context()
+85 |     execution_date = context["execution_date"]
+86 |     next_ds = context["next_ds"]
+   |                       ^^^^^^^^^ AIR302
+87 |     next_ds_nodash = context["next_ds_nodash"]
+88 |     next_execution_date = context["next_execution_date"]
+   |
+
+AIR302_context.py:87:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
+   |
+85 |     execution_date = context["execution_date"]
+86 |     next_ds = context["next_ds"]
+87 |     next_ds_nodash = context["next_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+88 |     next_execution_date = context["next_execution_date"]
+89 |     prev_ds = context["prev_ds"]
+   |
+
+AIR302_context.py:88:35: AIR302 `next_execution_date` is removed in Airflow 3.0
+   |
+86 |     next_ds = context["next_ds"]
+87 |     next_ds_nodash = context["next_ds_nodash"]
+88 |     next_execution_date = context["next_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+89 |     prev_ds = context["prev_ds"]
+90 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |
+
+AIR302_context.py:89:23: AIR302 `prev_ds` is removed in Airflow 3.0
+   |
+87 |     next_ds_nodash = context["next_ds_nodash"]
+88 |     next_execution_date = context["next_execution_date"]
+89 |     prev_ds = context["prev_ds"]
+   |                       ^^^^^^^^^ AIR302
+90 |     prev_ds_nodash = context["prev_ds_nodash"]
+91 |     prev_execution_date = context["prev_execution_date"]
+   |
+
+AIR302_context.py:90:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
+   |
+88 |     next_execution_date = context["next_execution_date"]
+89 |     prev_ds = context["prev_ds"]
+90 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+91 |     prev_execution_date = context["prev_execution_date"]
+92 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |
+
+AIR302_context.py:91:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
+   |
+89 |     prev_ds = context["prev_ds"]
+90 |     prev_ds_nodash = context["prev_ds_nodash"]
+91 |     prev_execution_date = context["prev_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+92 |     prev_execution_date_success = context["prev_execution_date_success"]
+93 |     tomorrow_ds = context["tomorrow_ds"]
+   |
+
+AIR302_context.py:92:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
+   |
+90 |     prev_ds_nodash = context["prev_ds_nodash"]
+91 |     prev_execution_date = context["prev_execution_date"]
+92 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
+93 |     tomorrow_ds = context["tomorrow_ds"]
+94 |     yesterday_ds = context["yesterday_ds"]
+   |
+
+AIR302_context.py:93:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
+   |
+91 |     prev_execution_date = context["prev_execution_date"]
+92 |     prev_execution_date_success = context["prev_execution_date_success"]
+93 |     tomorrow_ds = context["tomorrow_ds"]
+   |                           ^^^^^^^^^^^^^ AIR302
+94 |     yesterday_ds = context["yesterday_ds"]
+95 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:94:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
+   |
+92 |     prev_execution_date_success = context["prev_execution_date_success"]
+93 |     tomorrow_ds = context["tomorrow_ds"]
+94 |     yesterday_ds = context["yesterday_ds"]
+   |                            ^^^^^^^^^^^^^^ AIR302
+95 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:95:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
+   |
+93 |     tomorrow_ds = context["tomorrow_ds"]
+94 |     yesterday_ds = context["yesterday_ds"]
+95 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+96 |
+97 | class CustomOperator(BaseOperator):
+   |
+
+AIR302_context.py:112:45: AIR302 `execution_date` is removed in Airflow 3.0
     |
-134 | class CustomOperator(BaseOperator):
-135 |     def execute(self, next_ds, context):
-    |                       ^^^^^^^ AIR302
-136 |         execution_date = context["execution_date"]
-137 |         next_ds = context["next_ds"]
+111 | @task
+112 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
+    |                                             ^^^^^^^^^^^^^^ AIR302
+113 |     print("execution date", execution_date)
+114 |     print("access invalid key", context.get("conf"))
+    |
+
+AIR302_context.py:112:61: AIR302 `tomorrow_ds` is removed in Airflow 3.0
+    |
+111 | @task
+112 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
+    |                                                             ^^^^^^^^^^^ AIR302
+113 |     print("execution date", execution_date)
+114 |     print("access invalid key", context.get("conf"))
+    |
+
+AIR302_context.py:114:45: AIR302 `conf` is removed in Airflow 3.0
+    |
+112 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
+113 |     print("execution date", execution_date)
+114 |     print("access invalid key", context.get("conf"))
+    |                                             ^^^^^^ AIR302
+115 |
+116 | @task(task_id="print_the_context")
+    |
+
+AIR302_context.py:120:22: AIR302 `tomorrow_ds` is removed in Airflow 3.0
+    |
+118 |     """Print the Airflow context and ds variable from the context."""
+119 |     print(ds)
+120 |     print(kwargs.get("tomorrow_ds"))
+    |                      ^^^^^^^^^^^^^ AIR302
+121 |     c = get_current_context()
+122 |     c.get("execution_date")
+    |
+
+AIR302_context.py:122:11: AIR302 `execution_date` is removed in Airflow 3.0
+    |
+120 |     print(kwargs.get("tomorrow_ds"))
+121 |     c = get_current_context()
+122 |     c.get("execution_date")
+    |           ^^^^^^^^^^^^^^^^ AIR302
+123 |
+124 | class CustomOperatorNew(BaseOperator):
     |

scripts/knot_benchmark/src/benchmark/cases.py

@@ -68,7 +68,7 @@ class Knot(Tool):
         )
 
     def cold_command(self, project: Project, venv: Venv) -> Command:
-        command = [str(self.path), "-v"]
+        command = [str(self.path), "check", "-v"]
 
         assert len(project.include) < 2, "Knot doesn't support multiple source folders"