formatter

.github/actionlint.yaml

@@ -6,5 +6,4 @@ self-hosted-runner:
   labels:
     - depot-ubuntu-latest-8
     - depot-ubuntu-22.04-16
-    - github-windows-2025-x86_64-8
-    - github-windows-2025-x86_64-16
+    - windows-latest-xlarge

.github/workflows/ci.yaml

@@ -203,7 +203,7 @@ jobs:
 
   cargo-test-windows:
     name: "cargo test (windows)"
-    runs-on: github-windows-2025-x86_64-16
+    runs-on: windows-latest-xlarge
     needs: determine_changes
     if: ${{ needs.determine_changes.outputs.code == 'true' || github.ref == 'refs/heads/main' }}
     timeout-minutes: 20

Cargo.lock

@@ -2917,6 +2917,7 @@ dependencies = [
  "itertools 0.14.0",
  "memchr",
  "ruff_cache",
+ "ruff_index",
  "ruff_macros",
  "ruff_python_trivia",
  "ruff_source_file",

crates/red_knot_python_semantic/src/types.rs

@@ -963,27 +963,33 @@ impl<'db> Type<'db> {
             // `Literal[str]` is a subtype of `type` because the `str` class object is an instance of its metaclass `type`.
             // `Literal[abc.ABC]` is a subtype of `abc.ABCMeta` because the `abc.ABC` class object
             // is an instance of its metaclass `abc.ABCMeta`.
-            (Type::ClassLiteral(ClassLiteralType { class }), _) => class
-                .metaclass(db)
-                .to_instance(db)
-                .is_subtype_of(db, target),
+            (
+                Type::ClassLiteral(ClassLiteralType { class: self_class }),
+                Type::Instance(InstanceType {
+                    class: target_class,
+                }),
+            ) => self_class.is_instance_of(db, target_class),
 
             // `type[str]` (== `SubclassOf("str")` in red-knot) describes all possible runtime subclasses
             // of the class object `str`. It is a subtype of `type` (== `Instance("type")`) because `str`
             // is an instance of `type`, and so all possible subclasses of `str` will also be instances of `type`.
             //
             // Similarly `type[enum.Enum]`  is a subtype of `enum.EnumMeta` because `enum.Enum`
-            // is an instance of `enum.EnumMeta`. `type[Any]` and `type[Unknown]` do not participate in subtyping,
-            // however, as they are not fully static types.
-            (Type::SubclassOf(subclass_of_ty), _) => subclass_of_ty
+            // is an instance of `enum.EnumMeta`.
+            (
+                Type::SubclassOf(subclass_of_ty),
+                Type::Instance(InstanceType {
+                    class: target_class,
+                }),
+            ) => subclass_of_ty
                 .subclass_of()
                 .into_class()
-                .is_some_and(|class| {
-                    class
-                        .metaclass(db)
-                        .to_instance(db)
-                        .is_subtype_of(db, target)
-                }),
+                .is_some_and(|subclass_class| subclass_class.is_instance_of(db, target_class)),
+
+            // Other than the cases enumerated above, `type[]` and class-literal types just delegate to `Instance("type")`
+            (Type::SubclassOf(_) | Type::ClassLiteral(_), _) => {
+                KnownClass::Type.to_instance(db).is_subtype_of(db, target)
+            }
 
             // For example: `Type::KnownInstance(KnownInstanceType::Type)` is a subtype of `Type::Instance(_SpecialForm)`,
             // because `Type::KnownInstance(KnownInstanceType::Type)` is a set with exactly one runtime value in it
@@ -1376,16 +1382,14 @@ impl<'db> Type<'db> {
                 !KnownClass::Slice.is_subclass_of(db, class)
             }
 
-            // A class-literal type `X` is always disjoint from an instance type `Y`,
-            // unless the type expressing "all instances of `Z`" is a subtype of of `Y`,
-            // where `Z` is `X`'s metaclass.
-            (Type::ClassLiteral(ClassLiteralType { class }), instance @ Type::Instance(_))
-            | (instance @ Type::Instance(_), Type::ClassLiteral(ClassLiteralType { class })) => {
-                !class
-                    .metaclass(db)
-                    .to_instance(db)
-                    .is_subtype_of(db, instance)
-            }
+            (
+                Type::ClassLiteral(ClassLiteralType { class: class_a }),
+                Type::Instance(InstanceType { class: class_b }),
+            )
+            | (
+                Type::Instance(InstanceType { class: class_b }),
+                Type::ClassLiteral(ClassLiteralType { class: class_a }),
+            ) => !class_a.is_instance_of(db, class_b),
 
             (Type::FunctionLiteral(..), Type::Instance(InstanceType { class }))
             | (Type::Instance(InstanceType { class }), Type::FunctionLiteral(..)) => {
@@ -3853,6 +3857,17 @@ impl<'db> Class<'db> {
         self.iter_mro(db).contains(&ClassBase::Class(other))
     }
 
+    /// Return `true` if this class object is an instance of the class `other`.
+    ///
+    /// A class is an instance of its metaclass; consequently,
+    /// a class will only ever be an instance of another class
+    /// if its metaclass is a subclass of that other class.
+    fn is_instance_of(self, db: &'db dyn Db, other: Class<'db>) -> bool {
+        self.metaclass(db).into_class_literal().is_some_and(
+            |ClassLiteralType { class: metaclass }| metaclass.is_subclass_of(db, other),
+        )
+    }
+
     /// Return the explicit `metaclass` of this class, if one is defined.
     ///
     /// ## Note

crates/ruff_db/src/parsed.rs

@@ -2,7 +2,7 @@ use std::fmt::Formatter;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use ruff_python_ast::{ModModule, PySourceType};
+use ruff_python_ast::{ModModuleId, PySourceType};
 use ruff_python_parser::{parse_unchecked_source, Parsed};
 
 use crate::files::{File, FilePath};
@@ -43,24 +43,24 @@ pub fn parsed_module(db: &dyn Db, file: File) -> ParsedModule {
 /// Cheap cloneable wrapper around the parsed module.
 #[derive(Clone)]
 pub struct ParsedModule {
-    inner: Arc<Parsed<ModModule>>,
+    inner: Arc<Parsed<ModModuleId>>,
 }
 
 impl ParsedModule {
-    pub fn new(parsed: Parsed<ModModule>) -> Self {
+    pub fn new(parsed: Parsed<ModModuleId>) -> Self {
         Self {
             inner: Arc::new(parsed),
         }
     }
 
     /// Consumes `self` and returns the Arc storing the parsed module.
-    pub fn into_arc(self) -> Arc<Parsed<ModModule>> {
+    pub fn into_arc(self) -> Arc<Parsed<ModModuleId>> {
         self.inner
     }
 }
 
 impl Deref for ParsedModule {
-    type Target = Parsed<ModModule>;
+    type Target = Parsed<ModModuleId>;
 
     fn deref(&self) -> &Self::Target {
         &self.inner

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

@@ -1,119 +0,0 @@
-import pendulum
-from airflow.models import DAG
-from airflow.operators.dummy import DummyOperator
-from datetime import datetime
-from airflow.plugins_manager import AirflowPlugin
-from airflow.models.baseoperator import BaseOperator
-from airflow.decorators import dag, task
-from airflow.providers.standard.operators.python import PythonOperator
-from airflow.utils.context import get_current_context
-
-def access_invalid_key_in_context(**context):
-    print("access invalid key", context["conf"])
-
-@task
-def access_invalid_key_task_out_of_dag(**context):
-    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"))
-
-    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):
-    # This should not throw an error as logical_date is part of airflow context.
-    logical_date = context["logical_date"]
-
-    # Removed usage - should trigger violations
-    execution_date = context["execution_date"]
-    next_ds = context["next_ds"]
-    next_ds_nodash = context["next_ds_nodash"]
-    next_execution_date = context["next_execution_date"]
-    prev_ds = context["prev_ds"]
-    prev_ds_nodash = context["prev_ds_nodash"]
-    prev_execution_date = context["prev_execution_date"]
-    prev_execution_date_success = context["prev_execution_date_success"]
-    tomorrow_ds = context["tomorrow_ds"]
-    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():
-    context = get_current_context()
-    execution_date = context["execution_date"]
-    next_ds = context["next_ds"]
-    next_ds_nodash = context["next_ds_nodash"]
-    next_execution_date = context["next_execution_date"]
-    prev_ds = context["prev_ds"]
-    prev_ds_nodash = context["prev_ds_nodash"]
-    prev_execution_date = context["prev_execution_date"]
-    prev_execution_date_success = context["prev_execution_date_success"]
-    tomorrow_ds = context["tomorrow_ds"]
-    yesterday_ds = context["yesterday_ds"]
-    yesterday_ds_nodash = context["yesterday_ds_nodash"]
-
-class CustomOperator(BaseOperator):
-    def execute(self, context):
-        execution_date = context["execution_date"]
-        next_ds = context["next_ds"]
-        next_ds_nodash = context["next_ds_nodash"]
-        next_execution_date = context["next_execution_date"]
-        prev_ds = context["prev_ds"]
-        prev_ds_nodash = context["prev_ds_nodash"]
-        prev_execution_date = context["prev_execution_date"]
-        prev_execution_date_success = context["prev_execution_date_success"]
-        tomorrow_ds = context["tomorrow_ds"]
-        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")

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

@@ -175,9 +175,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::NonPEP646Unpack) {
                 pyupgrade::rules::use_pep646_unpack(checker, subscript);
             }
-            if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::removed_in_3(checker, expr);
-            }
+
             pandas_vet::rules::subscript(checker, value, expr);
         }
         Expr::Tuple(ast::ExprTuple {

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

@@ -376,9 +376,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::PytestParameterWithDefaultArgument) {
                 flake8_pytest_style::rules::parameter_with_default_argument(checker, function_def);
             }
-            if checker.enabled(Rule::Airflow3Removal) {
-                airflow::rules::removed_in_3_function_def(checker, function_def);
-            }
         }
         Stmt::Return(_) => {
             if checker.enabled(Rule::ReturnOutsideFunction) {

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

@@ -18,7 +18,6 @@ mod tests {
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_class_attribute.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_airflow_plugin.py"))]
-    #[test_case(Rule::Airflow3Removal, Path::new("AIR302_context.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR303.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

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

@@ -1,19 +1,17 @@
-use crate::checkers::ast::Checker;
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::helpers::map_callable;
-use ruff_python_ast::AnyParameterRef;
 use ruff_python_ast::{
     name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, ExprContext, ExprName,
-    ExprStringLiteral, ExprSubscript, Stmt, StmtClassDef, StmtFunctionDef,
+    StmtClassDef,
 };
 use ruff_python_semantic::analyze::typing;
 use ruff_python_semantic::Modules;
 use ruff_python_semantic::ScopeKind;
-use ruff_python_semantic::SemanticModel;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 
+use crate::checkers::ast::Checker;
+
 /// ## What it does
 /// Checks for uses of deprecated Airflow functions and values.
 ///
@@ -73,64 +71,6 @@ impl Violation for Airflow3Removal {
     }
 }
 
-const REMOVED_CONTEXT_KEYS: [&str; 12] = [
-    "conf",
-    "execution_date",
-    "next_ds",
-    "next_ds_nodash",
-    "next_execution_date",
-    "prev_ds",
-    "prev_ds_nodash",
-    "prev_execution_date",
-    "prev_execution_date_success",
-    "tomorrow_ds",
-    "yesterday_ds",
-    "yesterday_ds_nodash",
-];
-
-fn extract_name_from_slice(slice: &Expr) -> Option<String> {
-    match slice {
-        Expr::StringLiteral(ExprStringLiteral { value, .. }) => Some(value.to_string()),
-        _ => None,
-    }
-}
-
-/// Check if a subscript expression accesses a removed Airflow context variable.
-/// If a removed key is found, push a corresponding diagnostic.
-fn removed_context_variable(checker: &mut Checker, subscript: &ExprSubscript) {
-    let ExprSubscript { value, slice, .. } = subscript;
-
-    let is_context_arg = if let Expr::Name(ExprName { id, .. }) = &**value {
-        id.as_str() == "context" || id.as_str().starts_with("**")
-    } else {
-        false
-    };
-
-    let is_current_context =
-        if let Some(qualname) = typing::resolve_assignment(value, checker.semantic()) {
-            matches!(
-                qualname.segments(),
-                ["airflow", "utils", "context", "get_current_context"]
-            )
-        } else {
-            false
-        };
-
-    if is_context_arg || is_current_context {
-        if let Some(key) = extract_name_from_slice(slice) {
-            if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) {
-                checker.diagnostics.push(Diagnostic::new(
-                    Airflow3Removal {
-                        deprecated: key,
-                        replacement: Replacement::None,
-                    },
-                    slice.range(),
-                ));
-            }
-        }
-    }
-}
-
 /// AIR302
 pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
@@ -147,7 +87,6 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
                 check_call_arguments(checker, &qualname, arguments);
             };
             check_method(checker, call_expr);
-            check_context_get(checker, call_expr);
         }
         Expr::Attribute(attribute_expr @ ExprAttribute { attr, .. }) => {
             check_name(checker, expr, attr.range());
@@ -161,43 +100,10 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
                 }
             }
         }
-        Expr::Subscript(subscript_expr) => {
-            removed_context_variable(checker, subscript_expr);
-        }
         _ => {}
     }
 }
 
-/// AIR302
-pub(crate) fn removed_in_3_function_def(checker: &mut Checker, function_def: &StmtFunctionDef) {
-    if !checker.semantic().seen_module(Modules::AIRFLOW) {
-        return;
-    }
-
-    if !is_airflow_task(function_def, checker.semantic()) {
-        return;
-    }
-
-    for param in function_def
-        .parameters
-        .posonlyargs
-        .iter()
-        .chain(function_def.parameters.args.iter())
-        .chain(function_def.parameters.kwonlyargs.iter())
-    {
-        let param_name = param.parameter.name.as_str();
-        if REMOVED_CONTEXT_KEYS.contains(&param_name) {
-            checker.diagnostics.push(Diagnostic::new(
-                Airflow3Removal {
-                    deprecated: param_name.to_string(),
-                    replacement: Replacement::None,
-                },
-                param.parameter.name.range(),
-            ));
-        }
-    }
-}
-
 #[derive(Debug, Eq, PartialEq)]
 enum Replacement {
     None,
@@ -341,143 +247,6 @@ fn check_class_attribute(checker: &mut Checker, attribute_expr: &ExprAttribute)
     }
 }
 
-/// Finds the parameter definition for a given name expression in a function.
-fn find_parameter<'a>(
-    semantic: &'a SemanticModel,
-    name: &'a ExprName,
-) -> Option<AnyParameterRef<'a>> {
-    let binding_id = semantic.only_binding(name)?;
-    let binding = semantic.binding(binding_id);
-    let StmtFunctionDef { parameters, .. } = binding.statement(semantic)?.as_function_def_stmt()?;
-    parameters
-        .iter()
-        .find(|parameter| parameter.name().range() == binding.range())
-}
-
-/// Checks whether an Airflow 3.0–removed context key is used in a function decorated with `@task`.
-///
-/// Specifically, it flags two scenarios for task decorated function:
-/// 1. A removed context variable passed in as a function parameter name (e.g., `execution_date`).
-/// 2. A removed context key accessed via `context.get("...")`.
-///
-/// # Examples
-///
-/// **Removed key used in `context.get(...)`:**
-/// ```python
-/// from airflow.decorators import task
-///
-/// @task
-/// def my_task(**context):
-///     # 'conf' is removed in Airflow 3.0
-///     print(context.get("conf"))
-/// ```
-///
-/// **Removed context variable as a parameter:**
-/// ```python
-/// from airflow.decorators import task
-///
-/// @task
-/// def another_task(execution_date, **kwargs):
-///     # 'execution_date' is removed in Airflow 3.0
-///     pass
-/// ```
-///
-/// **Accessing multiple keys:**
-/// ```python
-/// from airflow.decorators import task
-///
-/// @task
-/// def more_keys(**context):
-///     # 'prev_ds' is also removed in Airflow 3.0
-///     print(context.get("prev_ds"))
-/// ```
-fn check_context_get(checker: &mut Checker, call_expr: &ExprCall) {
-    if !is_taskflow(checker) {
-        return;
-    }
-
-    let Expr::Attribute(ExprAttribute { value, attr, .. }) = &*call_expr.func else {
-        return;
-    };
-
-    let is_named_context = if let Expr::Name(name) = &**value {
-        if let Some(parameter) = find_parameter(checker.semantic(), name) {
-            matches!(parameter.name().as_str(), "context" | "kwargs")
-                || parameter.name().as_str().starts_with("**")
-        } else {
-            false
-        }
-    } else {
-        false
-    };
-
-    let is_assigned_from_gcc =
-        if let Some(qualname) = typing::resolve_assignment(value, checker.semantic()) {
-            matches!(
-                qualname.segments(),
-                ["airflow", "utils", "context", "get_current_context"]
-            )
-        } else {
-            false
-        };
-
-    if !(is_named_context || is_assigned_from_gcc) {
-        return;
-    }
-
-    if attr.as_str() != "get" {
-        return;
-    }
-
-    for removed_key in REMOVED_CONTEXT_KEYS {
-        if let Some(argument) = call_expr.arguments.find_argument_value(removed_key, 0) {
-            if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = argument {
-                if value == removed_key {
-                    checker.diagnostics.push(Diagnostic::new(
-                        Airflow3Removal {
-                            deprecated: removed_key.to_string(),
-                            replacement: Replacement::None,
-                        },
-                        argument.range(),
-                    ));
-                }
-            }
-        }
-    }
-}
-
-/// Check whether the function is decorated by @task
-///
-///
-/// Examples for the above patterns:
-/// ```python
-/// from airflow.decorators import task
-///
-///
-/// @task
-/// def access_invalid_key_task_out_of_dag(**context):
-///     print("access invalid key", context.get("conf"))
-/// ```
-fn is_taskflow(checker: &mut Checker) -> bool {
-    let mut parents = checker.semantic().current_statements();
-    if let Some(Stmt::FunctionDef(StmtFunctionDef { decorator_list, .. })) =
-        parents.find(|stmt| stmt.is_function_def_stmt())
-    {
-        for decorator in decorator_list {
-            if checker
-                .semantic()
-                .resolve_qualified_name(map_callable(&decorator.expression))
-                .is_some_and(|qualified_name| {
-                    matches!(qualified_name.segments(), ["airflow", "decorators", "task"])
-                })
-            {
-                return true;
-            }
-        }
-    }
-    false
-}
-
 /// Check whether a removed Airflow class method is called.
 ///
 /// For example:
@@ -1086,14 +855,3 @@ fn is_airflow_builtin_or_provider(segments: &[&str], module: &str, symbol_suffix
         _ => false,
     }
 }
-
-/// Returns `true` if the given function is decorated with `@airflow.decorators.task`.
-fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) -> bool {
-    function_def.decorator_list.iter().any(|decorator| {
-        semantic
-            .resolve_qualified_name(map_callable(&decorator.expression))
-            .is_some_and(|qualified_name| {
-                matches!(qualified_name.segments(), ["airflow", "decorators", "task"])
-            })
-    })
-}

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

@@ -1,455 +0,0 @@
----
-source: crates/ruff_linter/src/rules/airflow/mod.rs
-snapshot_kind: text
----
-AIR302_context.py:12:41: AIR302 `conf` is removed in Airflow 3.0
-   |
-11 | def access_invalid_key_in_context(**context):
-12 |     print("access invalid key", context["conf"])
-   |                                         ^^^^^^ AIR302
-13 |
-14 | @task
-   |
-
-AIR302_context.py:16:45: AIR302 `conf` is removed in Airflow 3.0
-   |
-14 | @task
-15 | def access_invalid_key_task_out_of_dag(**context):
-16 |     print("access invalid key", context.get("conf"))
-   |                                             ^^^^^^ AIR302
-17 |
-18 | @dag(
-   |
-
-AIR302_context.py:27:49: AIR302 `conf` is removed in Airflow 3.0
-   |
-25 |     @task()
-26 |     def access_invalid_key_task(**context):
-27 |         print("access invalid key", context.get("conf"))
-   |                                                 ^^^^^^ AIR302
-28 |
-29 |     task1 = PythonOperator(
-   |
-
-AIR302_context.py:44:30: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-43 |     # Removed usage - should trigger violations
-44 |     execution_date = context["execution_date"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-45 |     next_ds = context["next_ds"]
-46 |     next_ds_nodash = context["next_ds_nodash"]
-   |
-
-AIR302_context.py:45:23: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-43 |     # Removed usage - should trigger violations
-44 |     execution_date = context["execution_date"]
-45 |     next_ds = context["next_ds"]
-   |                       ^^^^^^^^^ AIR302
-46 |     next_ds_nodash = context["next_ds_nodash"]
-47 |     next_execution_date = context["next_execution_date"]
-   |
-
-AIR302_context.py:46:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
-   |
-44 |     execution_date = context["execution_date"]
-45 |     next_ds = context["next_ds"]
-46 |     next_ds_nodash = context["next_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-47 |     next_execution_date = context["next_execution_date"]
-48 |     prev_ds = context["prev_ds"]
-   |
-
-AIR302_context.py:47:35: AIR302 `next_execution_date` is removed in Airflow 3.0
-   |
-45 |     next_ds = context["next_ds"]
-46 |     next_ds_nodash = context["next_ds_nodash"]
-47 |     next_execution_date = context["next_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-48 |     prev_ds = context["prev_ds"]
-49 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |
-
-AIR302_context.py:48:23: AIR302 `prev_ds` is removed in Airflow 3.0
-   |
-46 |     next_ds_nodash = context["next_ds_nodash"]
-47 |     next_execution_date = context["next_execution_date"]
-48 |     prev_ds = context["prev_ds"]
-   |                       ^^^^^^^^^ AIR302
-49 |     prev_ds_nodash = context["prev_ds_nodash"]
-50 |     prev_execution_date = context["prev_execution_date"]
-   |
-
-AIR302_context.py:49:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
-   |
-47 |     next_execution_date = context["next_execution_date"]
-48 |     prev_ds = context["prev_ds"]
-49 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-50 |     prev_execution_date = context["prev_execution_date"]
-51 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |
-
-AIR302_context.py:50:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
-   |
-48 |     prev_ds = context["prev_ds"]
-49 |     prev_ds_nodash = context["prev_ds_nodash"]
-50 |     prev_execution_date = context["prev_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-51 |     prev_execution_date_success = context["prev_execution_date_success"]
-52 |     tomorrow_ds = context["tomorrow_ds"]
-   |
-
-AIR302_context.py:51:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
-   |
-49 |     prev_ds_nodash = context["prev_ds_nodash"]
-50 |     prev_execution_date = context["prev_execution_date"]
-51 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
-52 |     tomorrow_ds = context["tomorrow_ds"]
-53 |     yesterday_ds = context["yesterday_ds"]
-   |
-
-AIR302_context.py:52:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-50 |     prev_execution_date = context["prev_execution_date"]
-51 |     prev_execution_date_success = context["prev_execution_date_success"]
-52 |     tomorrow_ds = context["tomorrow_ds"]
-   |                           ^^^^^^^^^^^^^ AIR302
-53 |     yesterday_ds = context["yesterday_ds"]
-54 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:53:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
-   |
-51 |     prev_execution_date_success = context["prev_execution_date_success"]
-52 |     tomorrow_ds = context["tomorrow_ds"]
-53 |     yesterday_ds = context["yesterday_ds"]
-   |                            ^^^^^^^^^^^^^^ AIR302
-54 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:54:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
-   |
-52 |     tomorrow_ds = context["tomorrow_ds"]
-53 |     yesterday_ds = context["yesterday_ds"]
-54 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-55 |
-56 | with DAG(
-   |
-
-AIR302_context.py:58:5: AIR302 [*] `schedule_interval` is removed in Airflow 3.0
-   |
-56 | with DAG(
-57 |     dag_id="example_dag",
-58 |     schedule_interval="@daily",
-   |     ^^^^^^^^^^^^^^^^^ AIR302
-59 |     start_date=datetime(2023, 1, 1),
-60 |     template_searchpath=["/templates"],
-   |
-   = help: Use `schedule` instead
-
-ℹ Safe fix
-55 55 | 
-56 56 | with DAG(
-57 57 |     dag_id="example_dag",
-58    |-    schedule_interval="@daily",
-   58 |+    schedule="@daily",
-59 59 |     start_date=datetime(2023, 1, 1),
-60 60 |     template_searchpath=["/templates"],
-61 61 | ) as dag:
-
-AIR302_context.py:62:13: AIR302 `airflow.operators.dummy.DummyOperator` is removed in Airflow 3.0
-   |
-60 |     template_searchpath=["/templates"],
-61 | ) as dag:
-62 |     task1 = DummyOperator(
-   |             ^^^^^^^^^^^^^ AIR302
-63 |         task_id="task1",
-64 |         params={
-   |
-   = help: Use `airflow.operators.empty.EmptyOperator` instead
-
-AIR302_context.py:75:57: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-73 |     name = "custom_macros"
-74 |     macros = {
-75 |         "execution_date_macro": lambda context: context["execution_date"],
-   |                                                         ^^^^^^^^^^^^^^^^ AIR302
-76 |         "next_ds_macro": lambda context: context["next_ds"]
-77 |     }
-   |
-
-AIR302_context.py:76:50: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-74 |     macros = {
-75 |         "execution_date_macro": lambda context: context["execution_date"],
-76 |         "next_ds_macro": lambda context: context["next_ds"]
-   |                                                  ^^^^^^^^^ AIR302
-77 |     }
-   |
-
-AIR302_context.py:82:30: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-80 | def print_config():
-81 |     context = get_current_context()
-82 |     execution_date = context["execution_date"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-83 |     next_ds = context["next_ds"]
-84 |     next_ds_nodash = context["next_ds_nodash"]
-   |
-
-AIR302_context.py:83:23: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-81 |     context = get_current_context()
-82 |     execution_date = context["execution_date"]
-83 |     next_ds = context["next_ds"]
-   |                       ^^^^^^^^^ AIR302
-84 |     next_ds_nodash = context["next_ds_nodash"]
-85 |     next_execution_date = context["next_execution_date"]
-   |
-
-AIR302_context.py:84:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0
-   |
-82 |     execution_date = context["execution_date"]
-83 |     next_ds = context["next_ds"]
-84 |     next_ds_nodash = context["next_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-85 |     next_execution_date = context["next_execution_date"]
-86 |     prev_ds = context["prev_ds"]
-   |
-
-AIR302_context.py:85:35: AIR302 `next_execution_date` is removed in Airflow 3.0
-   |
-83 |     next_ds = context["next_ds"]
-84 |     next_ds_nodash = context["next_ds_nodash"]
-85 |     next_execution_date = context["next_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-86 |     prev_ds = context["prev_ds"]
-87 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |
-
-AIR302_context.py:86:23: AIR302 `prev_ds` is removed in Airflow 3.0
-   |
-84 |     next_ds_nodash = context["next_ds_nodash"]
-85 |     next_execution_date = context["next_execution_date"]
-86 |     prev_ds = context["prev_ds"]
-   |                       ^^^^^^^^^ AIR302
-87 |     prev_ds_nodash = context["prev_ds_nodash"]
-88 |     prev_execution_date = context["prev_execution_date"]
-   |
-
-AIR302_context.py:87:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
-   |
-85 |     next_execution_date = context["next_execution_date"]
-86 |     prev_ds = context["prev_ds"]
-87 |     prev_ds_nodash = context["prev_ds_nodash"]
-   |                              ^^^^^^^^^^^^^^^^ AIR302
-88 |     prev_execution_date = context["prev_execution_date"]
-89 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |
-
-AIR302_context.py:88:35: AIR302 `prev_execution_date` is removed in Airflow 3.0
-   |
-86 |     prev_ds = context["prev_ds"]
-87 |     prev_ds_nodash = context["prev_ds_nodash"]
-88 |     prev_execution_date = context["prev_execution_date"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-89 |     prev_execution_date_success = context["prev_execution_date_success"]
-90 |     tomorrow_ds = context["tomorrow_ds"]
-   |
-
-AIR302_context.py:89:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
-   |
-87 |     prev_ds_nodash = context["prev_ds_nodash"]
-88 |     prev_execution_date = context["prev_execution_date"]
-89 |     prev_execution_date_success = context["prev_execution_date_success"]
-   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
-90 |     tomorrow_ds = context["tomorrow_ds"]
-91 |     yesterday_ds = context["yesterday_ds"]
-   |
-
-AIR302_context.py:90:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-   |
-88 |     prev_execution_date = context["prev_execution_date"]
-89 |     prev_execution_date_success = context["prev_execution_date_success"]
-90 |     tomorrow_ds = context["tomorrow_ds"]
-   |                           ^^^^^^^^^^^^^ AIR302
-91 |     yesterday_ds = context["yesterday_ds"]
-92 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:91:28: AIR302 `yesterday_ds` is removed in Airflow 3.0
-   |
-89 |     prev_execution_date_success = context["prev_execution_date_success"]
-90 |     tomorrow_ds = context["tomorrow_ds"]
-91 |     yesterday_ds = context["yesterday_ds"]
-   |                            ^^^^^^^^^^^^^^ AIR302
-92 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |
-
-AIR302_context.py:92:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
-   |
-90 |     tomorrow_ds = context["tomorrow_ds"]
-91 |     yesterday_ds = context["yesterday_ds"]
-92 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
-   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
-93 |
-94 | class CustomOperator(BaseOperator):
-   |
-
-AIR302_context.py:96:34: AIR302 `execution_date` is removed in Airflow 3.0
-   |
-94 | class CustomOperator(BaseOperator):
-95 |     def execute(self, context):
-96 |         execution_date = context["execution_date"]
-   |                                  ^^^^^^^^^^^^^^^^ AIR302
-97 |         next_ds = context["next_ds"]
-98 |         next_ds_nodash = context["next_ds_nodash"]
-   |
-
-AIR302_context.py:97:27: AIR302 `next_ds` is removed in Airflow 3.0
-   |
-95 |     def execute(self, context):
-96 |         execution_date = context["execution_date"]
-97 |         next_ds = context["next_ds"]
-   |                           ^^^^^^^^^ AIR302
-98 |         next_ds_nodash = context["next_ds_nodash"]
-99 |         next_execution_date = context["next_execution_date"]
-   |
-
-AIR302_context.py:98:34: AIR302 `next_ds_nodash` is removed in Airflow 3.0
-    |
- 96 |         execution_date = context["execution_date"]
- 97 |         next_ds = context["next_ds"]
- 98 |         next_ds_nodash = context["next_ds_nodash"]
-    |                                  ^^^^^^^^^^^^^^^^ AIR302
- 99 |         next_execution_date = context["next_execution_date"]
-100 |         prev_ds = context["prev_ds"]
-    |
-
-AIR302_context.py:99:39: AIR302 `next_execution_date` is removed in Airflow 3.0
-    |
- 97 |         next_ds = context["next_ds"]
- 98 |         next_ds_nodash = context["next_ds_nodash"]
- 99 |         next_execution_date = context["next_execution_date"]
-    |                                       ^^^^^^^^^^^^^^^^^^^^^ AIR302
-100 |         prev_ds = context["prev_ds"]
-101 |         prev_ds_nodash = context["prev_ds_nodash"]
-    |
-
-AIR302_context.py:100:27: AIR302 `prev_ds` is removed in Airflow 3.0
-    |
- 98 |         next_ds_nodash = context["next_ds_nodash"]
- 99 |         next_execution_date = context["next_execution_date"]
-100 |         prev_ds = context["prev_ds"]
-    |                           ^^^^^^^^^ AIR302
-101 |         prev_ds_nodash = context["prev_ds_nodash"]
-102 |         prev_execution_date = context["prev_execution_date"]
-    |
-
-AIR302_context.py:101:34: AIR302 `prev_ds_nodash` is removed in Airflow 3.0
-    |
- 99 |         next_execution_date = context["next_execution_date"]
-100 |         prev_ds = context["prev_ds"]
-101 |         prev_ds_nodash = context["prev_ds_nodash"]
-    |                                  ^^^^^^^^^^^^^^^^ AIR302
-102 |         prev_execution_date = context["prev_execution_date"]
-103 |         prev_execution_date_success = context["prev_execution_date_success"]
-    |
-
-AIR302_context.py:102:39: AIR302 `prev_execution_date` is removed in Airflow 3.0
-    |
-100 |         prev_ds = context["prev_ds"]
-101 |         prev_ds_nodash = context["prev_ds_nodash"]
-102 |         prev_execution_date = context["prev_execution_date"]
-    |                                       ^^^^^^^^^^^^^^^^^^^^^ AIR302
-103 |         prev_execution_date_success = context["prev_execution_date_success"]
-104 |         tomorrow_ds = context["tomorrow_ds"]
-    |
-
-AIR302_context.py:103:47: AIR302 `prev_execution_date_success` is removed in Airflow 3.0
-    |
-101 |         prev_ds_nodash = context["prev_ds_nodash"]
-102 |         prev_execution_date = context["prev_execution_date"]
-103 |         prev_execution_date_success = context["prev_execution_date_success"]
-    |                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
-104 |         tomorrow_ds = context["tomorrow_ds"]
-105 |         yesterday_ds = context["yesterday_ds"]
-    |
-
-AIR302_context.py:104:31: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-    |
-102 |         prev_execution_date = context["prev_execution_date"]
-103 |         prev_execution_date_success = context["prev_execution_date_success"]
-104 |         tomorrow_ds = context["tomorrow_ds"]
-    |                               ^^^^^^^^^^^^^ AIR302
-105 |         yesterday_ds = context["yesterday_ds"]
-106 |         yesterday_ds_nodash = context["yesterday_ds_nodash"]
-    |
-
-AIR302_context.py:105:32: AIR302 `yesterday_ds` is removed in Airflow 3.0
-    |
-103 |         prev_execution_date_success = context["prev_execution_date_success"]
-104 |         tomorrow_ds = context["tomorrow_ds"]
-105 |         yesterday_ds = context["yesterday_ds"]
-    |                                ^^^^^^^^^^^^^^ AIR302
-106 |         yesterday_ds_nodash = context["yesterday_ds_nodash"]
-    |
-
-AIR302_context.py:106:39: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0
-    |
-104 |         tomorrow_ds = context["tomorrow_ds"]
-105 |         yesterday_ds = context["yesterday_ds"]
-106 |         yesterday_ds_nodash = context["yesterday_ds_nodash"]
-    |                                       ^^^^^^^^^^^^^^^^^^^^^ AIR302
-107 |
-108 | @task
-    |
-
-AIR302_context.py:109:45: AIR302 `execution_date` is removed in Airflow 3.0
-    |
-108 | @task
-109 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
-    |                                             ^^^^^^^^^^^^^^ AIR302
-110 |     print("execution date", execution_date)
-111 |     print("access invalid key", context.get("conf"))
-    |
-
-AIR302_context.py:109:61: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-    |
-108 | @task
-109 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
-    |                                                             ^^^^^^^^^^^ AIR302
-110 |     print("execution date", execution_date)
-111 |     print("access invalid key", context.get("conf"))
-    |
-
-AIR302_context.py:111:45: AIR302 `conf` is removed in Airflow 3.0
-    |
-109 | def access_invalid_argument_task_out_of_dag(execution_date, tomorrow_ds, logical_date, **context):
-110 |     print("execution date", execution_date)
-111 |     print("access invalid key", context.get("conf"))
-    |                                             ^^^^^^ AIR302
-112 |
-113 | @task(task_id="print_the_context")
-    |
-
-AIR302_context.py:117:22: AIR302 `tomorrow_ds` is removed in Airflow 3.0
-    |
-115 |     """Print the Airflow context and ds variable from the context."""
-116 |     print(ds)
-117 |     print(kwargs.get("tomorrow_ds"))
-    |                      ^^^^^^^^^^^^^ AIR302
-118 |     c = get_current_context()
-119 |     c.get("execution_date")
-    |
-
-AIR302_context.py:119:11: AIR302 `execution_date` is removed in Airflow 3.0
-    |
-117 |     print(kwargs.get("tomorrow_ds"))
-118 |     c = get_current_context()
-119 |     c.get("execution_date")
-    |           ^^^^^^^^^^^^^^^^ AIR302
-    |

crates/ruff_macros/src/ast_node.rs

@@ -0,0 +1,321 @@
+use heck::ToSnakeCase;
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::spanned::Spanned;
+use syn::{Attribute, Error, Fields, Ident, ItemEnum, Result, Type, Variant};
+
+pub(crate) fn generate_ast_enum(input: ItemEnum) -> Result<TokenStream> {
+    let ast_enum = AstEnum::new(input)?;
+    let id_enum = generate_id_enum(&ast_enum);
+    let node_enum = generate_node_enum(&ast_enum);
+    let node_enum_node_method = generate_node_enum_node_method(&ast_enum);
+    let node_enum_ranged_impl = generate_node_enum_ranged_impl(&ast_enum);
+    let variant_ids = generate_variant_ids(&ast_enum);
+    let storage = generate_storage(&ast_enum);
+    Ok(quote! {
+        #id_enum
+        #node_enum
+        #node_enum_node_method
+        #node_enum_ranged_impl
+        #variant_ids
+        #storage
+    })
+}
+
+fn snake_case(node_ident: &Ident) -> Ident {
+    let node_string = node_ident.to_string().to_snake_case();
+    Ident::new(&node_string, node_ident.span())
+}
+
+fn concat(prefix: &str, id: &Ident, suffix: &str) -> Ident {
+    let mut id_string = id.to_string();
+    id_string.insert_str(0, prefix);
+    id_string.push_str(suffix);
+    Ident::new(&id_string, id.span())
+}
+
+/// Describes one of the enums that holds syntax nodes (e.g. Mod, Stmt)
+struct AstEnum {
+    /// The base name of the enums (e.g. Mod, Stmt)
+    base_enum_name: Ident,
+    /// The syntax node variants for this enum
+    variants: Vec<AstVariant>,
+}
+
+/// Describes one specific syntax node (e.g. ModExpression, StmtIf)
+struct AstVariant {
+    /// The name of the variant within its containing enum (e.g. Expression, If)
+    variant_name: Ident,
+    /// The struct type defining the contents of this syntax node (e.g. ModExpression, StmtIf)
+    node_ty: Ident,
+    /// All of the attributes attached to this variant
+    attrs: Vec<Attribute>,
+}
+
+impl AstEnum {
+    fn new(input: ItemEnum) -> Result<AstEnum> {
+        let base_enum_name = input.ident;
+        let variants: Result<Vec<_>> = input.variants.into_iter().map(AstVariant::new).collect();
+        let variants = variants?;
+        Ok(AstEnum {
+            base_enum_name,
+            variants,
+        })
+    }
+
+    fn map_variants<'a, B, F>(&'a self, f: F) -> impl Iterator<Item = B> + 'a
+    where
+        F: FnMut(&AstVariant) -> B + 'a,
+    {
+        self.variants.iter().map(f)
+    }
+
+    /// The name of the enum containing syntax node IDs (e.g. ModId, StmtId)
+    fn id_enum_ty(&self) -> Ident {
+        concat("", &self.base_enum_name, "Id")
+    }
+
+    /// The name of the enum containing references to syntax nodes (e.g. ModRef, StmtRef)
+    fn ref_enum_ty(&self) -> Ident {
+        concat("", &self.base_enum_name, "Ref")
+    }
+
+    /// The name of the storage type for this enum (e.g. ModStorage)
+    fn enum_storage_ty(&self) -> Ident {
+        concat("", &self.base_enum_name, "Storage")
+    }
+
+    /// The name of the storage field in Ast (e.g. mod_storage)
+    fn enum_storage_field(&self) -> Ident {
+        snake_case(&self.enum_storage_ty())
+    }
+}
+
+impl AstVariant {
+    fn new(variant: Variant) -> Result<AstVariant> {
+        let Fields::Unnamed(fields) = &variant.fields else {
+            return Err(Error::new(
+                variant.fields.span(),
+                "Each AstNode variant must have a single unnamed field",
+            ));
+        };
+        let mut fields = fields.unnamed.iter();
+        let field = fields.next().ok_or_else(|| {
+            Error::new(
+                variant.fields.span(),
+                "Each AstNode variant must have a single unnamed field",
+            )
+        })?;
+        if fields.next().is_some() {
+            return Err(Error::new(
+                variant.fields.span(),
+                "Each AstNode variant must have a single unnamed field",
+            ));
+        }
+        let Type::Path(field_ty) = &field.ty else {
+            return Err(Error::new(
+                field.ty.span(),
+                "Each AstNode variant must wrap a simple Id type",
+            ));
+        };
+        let node_ty = field_ty.path.require_ident()?.clone();
+        Ok(AstVariant {
+            variant_name: variant.ident,
+            node_ty,
+            attrs: variant.attrs,
+        })
+    }
+
+    /// The name of the ID type for this variant's syntax node (e.g. ModExpressionId, StmtIfId)
+    fn id_ty(&self) -> Ident {
+        concat("", &self.node_ty, "Id")
+    }
+
+    /// The name of the storage field in the containing enum storage type (e.g.
+    /// mod_expression_storage)
+    fn variant_storage_field(&self) -> Ident {
+        concat("", &snake_case(&self.node_ty), "_storage")
+    }
+
+    /// The name of the method that adds a new syntax node to an [Ast] (e.g. `add_mod_expression`)
+}
+
+/// Generates the enum containing syntax node IDs (e.g. ModId, StmtId)
+fn generate_id_enum(ast_enum: &AstEnum) -> TokenStream {
+    let id_enum_ty = ast_enum.id_enum_ty();
+    let id_enum_variants = ast_enum.map_variants(|v| {
+        let AstVariant {
+            variant_name,
+            attrs,
+            ..
+        } = v;
+        let id_ty = v.id_ty();
+        quote! {
+            #( #attrs )*
+            #variant_name(#id_ty)
+        }
+    });
+    quote! {
+        #[automatically_derived]
+        #[derive(Copy, Clone, Debug, PartialEq, is_macro::Is)]
+        pub enum #id_enum_ty {
+            #( #id_enum_variants ),*
+        }
+    }
+}
+
+fn generate_ref_enum(ast_enum: &AstEnum) -> TokenStream {
+    let ref_enum_ty = ast_enum.ref_enum_ty();
+    let variants = ast_enum.map_variants(|v| {
+        let AstVariant {
+            attrs,
+            variant_name,
+            node_ty,
+            ..
+        } = v;
+        quote! {
+            #( #attrs )*
+            #variant_name(crate::Node<'a, &'a #node_ty>)
+        }
+    });
+    quote! {
+        #[automatically_derived]
+        #[derive(Copy, Clone, Debug, PartialEq, is_macro::Is)]
+        pub enum #node_ident<'a> {
+            #( #variants ),*
+        }
+    }
+}
+
+fn generate_node_enum_node_method(ast_enum: &AstEnum) -> TokenStream {
+    let id_enum_ty = ast_enum.id_enum_ty();
+    let ref_enum_ty = ast_enum.ref_enum_ty();
+    let variants = ast_enum.map_variants(|v| {
+        let AstVariant { variant_name, .. } = v;
+        quote! { #id_enum_ty::#variant_name(id) => #ref_enum_ty::#variant_name(self.ast.wrap(&self.ast[id])) }
+    });
+    quote! {
+        #[automatically_derived]
+        impl<'a> crate::Node<'a, #id_enum_ty> {
+            #[inline]
+            pub fn node(&self) -> #ref_enum_ty<'a> {
+                match self.node {
+                    #( #variants ),*
+                }
+            }
+        }
+    }
+}
+
+fn generate_node_enum_ranged_impl(ast_enum: &AstEnum) -> TokenStream {
+    let ref_enum_ty = ast_enum.ref_enum_ty();
+    let variants = ast_enum.map_variants(|v| {
+        let AstVariant { variant_name, .. } = v;
+        quote! { #ref_enum_ty::#variant_name(node) => node.range() }
+    });
+    quote! {
+        #[automatically_derived]
+        impl ruff_text_size::Ranged for #ref_enum_ty<'_> {
+            fn range(&self) -> ruff_text_size::TextRange {
+                match self {
+                    #( #variants ),*
+                }
+            }
+        }
+    }
+}
+
+/// Generates the ID type for each syntax node struct in this enum.
+///
+/// We also define:
+///   - [Index] and [IndexMut] impls so that you can index into an [Ast] using the ID type
+///   - a `node` method on e.g. `Node<StmtIfId>` that returns a `Node<&StmtIf>`
+///   - [Ranged] impls for the `StmtIf` and `Node<&StmtIf>`
+fn generate_ids(ast_enum: &AstEnum) -> TokenStream {
+    let id_enum_ty = ast_enum.id_enum_ty();
+    let enum_storage_field = ast_enum.enum_storage_field();
+    let variants = ast_enum.map_variants(|v| {
+        let AstVariant { node_ty, .. } = v;
+        let id_ty = v.id_ty();
+        let variant_storage_field = v.variant_storage_field();
+        quote! {
+            #[automatically_derived]
+            #[ruff_index::newtype_index]
+            pub struct #id_ty;
+
+            #[automatically_derived]
+            impl std::ops::Index<#id_ty> for crate::Ast {
+                type Output = #node_ty;
+                #[inline]
+                fn index(&self, id: #id_ty) -> &#node_ty {
+                    &self.#enum_storage_field.#variant_storage_field[id]
+                }
+            }
+
+            #[automatically_derived]
+            impl std::ops::IndexMut<#id_ty> for crate::Ast {
+                #[inline]
+                fn index_mut(&mut self, id: #id_ty) -> &mut #node_ty {
+                    &mut self.#enum_storage_field.#variant_storage_field[id]
+                }
+            }
+
+            #[automatically_derived]
+            impl<'a> crate::Node<'a, #id_ty> {
+                #[inline]
+                pub fn node(&self) -> crate::Node<'a, &'a #node_ty> {
+                    self.ast.wrap(&self.ast[self.node])
+                }
+            }
+
+            #[automatically_derived]
+            impl<'a> ruff_text_size::Ranged for #node_ty {
+                fn range(&self) -> TextRange {
+                    self.range
+                }
+            }
+
+            #[automatically_derived]
+            impl<'a> ruff_text_size::Ranged for crate::Node<'a, &'a #node_ty> {
+                fn range(&self) -> TextRange {
+                    self.as_ref().range()
+                }
+            }
+        }
+    });
+    quote! { #( #variants )* }
+}
+
+fn generate_storage(ast_enum: &AstEnum) -> TokenStream {
+    let id_enum_ty = ast_enum.id_enum_ty();
+    let enum_storage_ty = ast_enum.enum_storage_ty();
+    let enum_storage_field = ast_enum.enum_storage_field();
+    let storage_fields = ast_enum.map_variants(|v| {
+        let AstVariant { id_ty, node_ty, .. } = v;
+        let variant_storage_field = v.variant_storage_field();
+        quote! { #variant_storage_field: ruff_index::IndexVec<#id_ty, #node_ty> }
+    });
+    let add_methods = ast_enum.map_variants(|v| {
+        let AstVariant { node_ty, .. } = v;
+        let variant_storage_field = v.variant_storage_field();
+        let method_name = concat("add_", vec_name, "");
+        quote! {
+            #[automatically_derived]
+            impl crate::Ast {
+                pub fn #method_name(&mut self, payload: #node_ty) -> #id_ident {
+                    #id_ident::#variant_name(self.#storage_field.#vec_name.push(payload))
+                }
+            }
+        }
+    });
+    quote! {
+        #[automatically_derived]
+        #[allow(clippy::derive_partial_eq_without_eq)]
+        #[derive(Clone, Default, PartialEq)]
+        pub(crate) struct #storage_ty {
+            #( #storage_fields ),*
+        }
+
+        #( #add_methods )*
+    }
+}

crates/ruff_python_ast/Cargo.toml

@@ -14,6 +14,7 @@ license = { workspace = true }
 
 [dependencies]
 ruff_cache = { workspace = true, optional = true }
+ruff_index = { workspace = true }
 ruff_macros = { workspace = true, optional = true }
 ruff_python_trivia = { workspace = true }
 ruff_source_file = { workspace = true }

crates/ruff_python_ast/generate.py

@@ -67,6 +67,7 @@ class Ast:
 class Group:
     name: str
     nodes: list[Node]
+    id_enum_ty: str
     owned_enum_ty: str
 
     add_suffix_to_is_methods: bool
@@ -75,6 +76,7 @@ class Group:
 
     def __init__(self, group_name: str, group: dict[str, Any]) -> None:
         self.name = group_name
+        self.id_enum_ty = group_name + "Id"
         self.owned_enum_ty = group_name
         self.ref_enum_ty = group_name + "Ref"
         self.add_suffix_to_is_methods = group.get("add_suffix_to_is_methods", False)
@@ -89,12 +91,16 @@ class Group:
 class Node:
     name: str
     variant: str
+    id_ty: str
     ty: str
+    storage_field: str
 
     def __init__(self, group: Group, node_name: str, node: dict[str, Any]) -> None:
         self.name = node_name
         self.variant = node.get("variant", node_name.removeprefix(group.name))
+        self.id_ty = node_name + "Id"
         self.ty = f"crate::{node_name}"
+        self.storage_field = to_snake_case(node_name)
 
 
 # ------------------------------------------------------------------------------
@@ -108,6 +114,75 @@ def write_preamble(out: list[str]) -> None:
     """)
 
 
+# ------------------------------------------------------------------------------
+# ID enum
+
+
+def write_ids(out: list[str], ast: Ast) -> None:
+    """
+    Create an ID type for each syntax node, and a per-group enum that contains a
+    syntax node ID.
+
+    ```rust
+    #[newindex_type]
+    pub struct TypeParamTypeVarId;
+    #[newindex_type]
+    pub struct TypeParamTypeVarTuple;
+    ...
+
+    pub enum TypeParamId {
+        TypeVar(TypeParamTypeVarId),
+        TypeVarTuple(TypeParamTypeVarTupleId),
+        ...
+    }
+    ```
+
+    Also creates:
+    - `impl From<TypeParamTypeVarId> for TypeParamId`
+    - `impl Ranged for TypeParamTypeVar`
+    - `fn TypeParamId::is_type_var() -> bool`
+
+    If the `add_suffix_to_is_methods` group option is true, then the
+    `is_type_var` method will be named `is_type_var_type_param`.
+    """
+
+    for node in ast.all_nodes:
+        out.append("")
+        out.append("#[ruff_index::newtype_index]")
+        out.append(f"pub struct {node.id_ty};")
+
+        out.append(f"""
+            impl ruff_text_size::Ranged for {node.ty} {{
+                fn range(&self) -> ruff_text_size::TextRange {{
+                    self.range
+                }}
+            }}
+        """)
+
+    for group in ast.groups:
+        out.append("")
+        if group.rustdoc is not None:
+            out.append(group.rustdoc)
+        out.append("#[derive(Clone, Copy, Debug, PartialEq, is_macro::Is)]")
+        out.append(f"pub enum {group.id_enum_ty} {{")
+        for node in group.nodes:
+            if group.add_suffix_to_is_methods:
+                is_name = to_snake_case(node.variant + group.name)
+                out.append(f'#[is(name = "{is_name}")]')
+            out.append(f"{node.variant}({node.id_ty}),")
+        out.append("}")
+
+        for node in group.nodes:
+            out.append(f"""
+            impl From<{node.id_ty}> for {group.id_enum_ty} {{
+                fn from(id: {node.id_ty}) -> Self {{
+                    Self::{node.variant}(id)
+                }}
+            }}
+            """)
+
+
+
 # ------------------------------------------------------------------------------
 # Owned enum
 
@@ -126,7 +201,6 @@ def write_owned_enum(out: list[str], ast: Ast) -> None:
 
     Also creates:
     - `impl Ranged for TypeParam`
-    - `TypeParam::visit_source_order`
     - `impl From<TypeParamTypeVar> for TypeParam`
     - `impl Ranged for TypeParamTypeVar`
     - `fn TypeParam::is_type_var() -> bool`
@@ -170,15 +244,6 @@ def write_owned_enum(out: list[str], ast: Ast) -> None:
         }
         """)
 
-    for node in ast.all_nodes:
-        out.append(f"""
-            impl ruff_text_size::Ranged for {node.ty} {{
-                fn range(&self) -> ruff_text_size::TextRange {{
-                    self.range
-                }}
-            }}
-        """)
-
     for group in ast.groups:
         out.append(f"""
             impl {group.owned_enum_ty} {{
@@ -210,17 +275,18 @@ def write_ref_enum(out: list[str], ast: Ast) -> None:
 
     ```rust
     pub enum TypeParamRef<'a> {
-        TypeVar(&'a TypeParamTypeVar),
-        TypeVarTuple(&'a TypeParamTypeVarTuple),
+        TypeVar(Node<'a, &'a TypeParamTypeVar>),
+        TypeVarTuple(Node<'a, &'a TypeParamTypeVarTuple>),
         ...
     }
     ```
 
     Also creates:
-    - `impl<'a> From<&'a TypeParam> for TypeParamRef<'a>`
-    - `impl<'a> From<&'a TypeParamTypeVar> for TypeParamRef<'a>`
+    - `impl<'a> From<Node<'a, &'a TypeParam>> for TypeParamRef<'a>`
+    - `impl<'a> From<Node<'a, &'a TypeParamTypeVar>> for TypeParamRef<'a>`
     - `impl Ranged for TypeParamRef<'_>`
     - `fn TypeParamRef::is_type_var() -> bool`
+    - `TypeParamRef::visit_source_order`
 
     The name of each variant can be customized via the `variant` node option. If
     the `add_suffix_to_is_methods` group option is true, then the `is_type_var`
@@ -237,17 +303,17 @@ def write_ref_enum(out: list[str], ast: Ast) -> None:
             if group.add_suffix_to_is_methods:
                 is_name = to_snake_case(node.variant + group.name)
                 out.append(f'#[is(name = "{is_name}")]')
-            out.append(f"""{node.variant}(&'a {node.ty}),""")
+            out.append(f"""{node.variant}(crate::Node<'a, &'a {node.ty}>),""")
         out.append("}")
 
         out.append(f"""
-            impl<'a> From<&'a {group.owned_enum_ty}> for {group.ref_enum_ty}<'a> {{
-                fn from(node: &'a {group.owned_enum_ty}) -> Self {{
-                    match node {{
+            impl<'a> From<crate::Node<'a, &'a {group.owned_enum_ty}>> for {group.ref_enum_ty}<'a> {{
+                fn from(node: crate::Node<'a, &'a {group.owned_enum_ty}>) -> Self {{
+                    match node.node {{
         """)
         for node in group.nodes:
             out.append(
-                f"{group.owned_enum_ty}::{node.variant}(node) => {group.ref_enum_ty}::{node.variant}(node),"
+                f"""{group.owned_enum_ty}::{node.variant}(n) => {group.ref_enum_ty}::{node.variant}(node.ast.wrap(n)),"""
             )
         out.append("""
                     }
@@ -257,8 +323,8 @@ def write_ref_enum(out: list[str], ast: Ast) -> None:
 
         for node in group.nodes:
             out.append(f"""
-            impl<'a> From<&'a {node.ty}> for {group.ref_enum_ty}<'a> {{
-                fn from(node: &'a {node.ty}) -> Self {{
+            impl<'a> From<crate::Node<'a, &'a {node.ty}>> for {group.ref_enum_ty}<'a> {{
+                fn from(node: crate::Node<'a, &'a {node.ty}>) -> Self {{
                     Self::{node.variant}(node)
                 }}
             }}
@@ -277,6 +343,112 @@ def write_ref_enum(out: list[str], ast: Ast) -> None:
         }
         """)
 
+    for group in ast.groups:
+        out.append(f"""
+            impl<'a> {group.ref_enum_ty}<'a> {{
+                #[allow(unused)]
+                pub(crate) fn visit_source_order<V>(self, visitor: &mut V)
+                where
+                    V: crate::visitor::source_order::SourceOrderVisitor<'a> + ?Sized,
+                {{
+                    match self {{
+        """)
+        for node in group.nodes:
+            out.append(
+                f"""{group.ref_enum_ty}::{node.variant}(node) => node.visit_source_order(visitor),"""
+            )
+        out.append("""
+                    }
+                }
+            }
+        """)
+
+
+# ------------------------------------------------------------------------------
+# AST storage
+
+
+def write_storage(out: list[str], ast: Ast) -> None:
+    """
+    Create the storage struct for all of the syntax nodes.
+
+    ```rust
+    pub(crate) struct Storage {
+        ...
+        pub(crate) type_param_type_var_id: IndexVec<TypeParamTypeVarId, TypeParamTypeVar>,
+        pub(crate) type_param_type_var_tuple_id: IndexVec<TypeParamTypeVarTupleId, TypeParamTypeVarTuple>,
+        ...
+    }
+    ```
+
+    Also creates:
+    - `impl AstId for TypeParamTypeVarId for Ast`
+    - `impl AstIdMut for TypeParamTypeVarId for Ast`
+    """
+
+    out.append("")
+    out.append("#[derive(Clone, Default, PartialEq)]")
+    out.append("pub(crate) struct Storage {")
+    for node in ast.all_nodes:
+        out.append(f"""pub(crate) {node.storage_field}: ruff_index::IndexVec<{node.id_ty}, {node.ty}>,""")
+    out.append("}")
+
+    for node in ast.all_nodes:
+        out.append(f"""
+            impl crate::ast::AstId for {node.id_ty} {{
+                type Output<'a> = crate::Node<'a, &'a {node.ty}>;
+                #[inline]
+                fn node<'a>(self, ast: &'a crate::Ast) -> Self::Output<'a> {{
+                    ast.wrap(&ast.storage.{node.storage_field}[self])
+                }}
+            }}
+        """)
+
+        out.append(f"""
+            impl crate::ast::AstIdMut for {node.id_ty} {{
+                type Output<'a> = crate::Node<'a, &'a mut {node.ty}>;
+                #[inline]
+                fn node_mut<'a>(self, ast: &'a mut crate::Ast) -> Self::Output<'a> {{
+                    ast.wrap(&mut ast.storage.{node.storage_field}[self])
+                }}
+            }}
+        """)
+
+        out.append(f"""
+            impl<'a> crate::Node<'a, {node.id_ty}> {{
+                #[inline]
+                pub fn node(self) -> crate::Node<'a, &'a {node.ty}> {{
+                    self.ast.node(self.node)
+                }}
+            }}
+        """)
+
+    for group in ast.groups:
+        out.append(f"""
+            impl crate::ast::AstId for {group.id_enum_ty} {{
+                type Output<'a> = {group.ref_enum_ty}<'a>;
+                #[inline]
+                fn node<'a>(self, ast: &'a crate::Ast) -> Self::Output<'a> {{
+                    match self {{
+        """)
+        for node in group.nodes:
+            out.append(f"""{group.id_enum_ty}::{node.variant}(node) => {group.ref_enum_ty}::{node.variant}(ast.node(node)),""")
+        out.append(f"""
+                    }}
+                }}
+            }}
+        """)
+
+        out.append(f"""
+            impl<'a> crate::Node<'a, {group.id_enum_ty}> {{
+                #[inline]
+                pub fn node(self) -> crate::Node<'a, &'a {node.ty}> {{
+                    self.ast.node(self.node)
+                }}
+            }}
+        """)
+
+
 
 # ------------------------------------------------------------------------------
 # AnyNodeRef
@@ -289,16 +461,16 @@ def write_anynoderef(out: list[str], ast: Ast) -> None:
     ```rust
     pub enum AnyNodeRef<'a> {
         ...
-        TypeParamTypeVar(&'a TypeParamTypeVar),
-        TypeParamTypeVarTuple(&'a TypeParamTypeVarTuple),
+        TypeParamTypeVar(Node<'a, &'a TypeParamTypeVar>),
+        TypeParamTypeVarTuple(Node<'a, &'a TypeParamTypeVarTuple>),
         ...
     }
     ```
 
     Also creates:
-    - `impl<'a> From<&'a TypeParam> for AnyNodeRef<'a>`
     - `impl<'a> From<TypeParamRef<'a>> for AnyNodeRef<'a>`
-    - `impl<'a> From<&'a TypeParamTypeVarTuple> for AnyNodeRef<'a>`
+    - `impl<'a> From<Node<'a, &'a TypeParam>> for AnyNodeRef<'a>`
+    - `impl<'a> From<Node<'a, &'a TypeParamTypeVarTuple>> for AnyNodeRef<'a>`
     - `impl Ranged for AnyNodeRef<'_>`
     - `fn AnyNodeRef::as_ptr(&self) -> std::ptr::NonNull<()>`
     - `fn AnyNodeRef::visit_preorder(self, visitor &mut impl SourceOrderVisitor)`
@@ -309,20 +481,20 @@ def write_anynoderef(out: list[str], ast: Ast) -> None:
     pub enum AnyNodeRef<'a> {
     """)
     for node in ast.all_nodes:
-        out.append(f"""{node.name}(&'a {node.ty}),""")
+        out.append(f"""{node.name}(crate::Node<'a, &'a {node.ty}>),""")
     out.append("""
     }
     """)
 
     for group in ast.groups:
         out.append(f"""
-            impl<'a> From<&'a {group.owned_enum_ty}> for AnyNodeRef<'a> {{
-                fn from(node: &'a {group.owned_enum_ty}) -> AnyNodeRef<'a> {{
-                    match node {{
+            impl<'a> From<crate::Node<'a, &'a {group.owned_enum_ty}>> for AnyNodeRef<'a> {{
+                fn from(node: crate::Node<'a, &'a {group.owned_enum_ty}>) -> AnyNodeRef<'a> {{
+                    match node.node {{
         """)
         for node in group.nodes:
             out.append(
-                f"{group.owned_enum_ty}::{node.variant}(node) => AnyNodeRef::{node.name}(node),"
+                f"{group.owned_enum_ty}::{node.variant}(n) => AnyNodeRef::{node.name}(node.ast.wrap(n)),"
             )
         out.append("""
                     }
@@ -347,8 +519,8 @@ def write_anynoderef(out: list[str], ast: Ast) -> None:
 
     for node in ast.all_nodes:
         out.append(f"""
-            impl<'a> From<&'a {node.ty}> for AnyNodeRef<'a> {{
-                fn from(node: &'a {node.ty}) -> AnyNodeRef<'a> {{
+            impl<'a> From<crate::Node<'a, &'a {node.ty}>> for AnyNodeRef<'a> {{
+                fn from(node: crate::Node<'a, &'a {node.ty}>) -> AnyNodeRef<'a> {{
                     AnyNodeRef::{node.name}(node)
                 }}
             }}
@@ -374,7 +546,7 @@ def write_anynoderef(out: list[str], ast: Ast) -> None:
     """)
     for node in ast.all_nodes:
         out.append(
-            f"AnyNodeRef::{node.name}(node) => std::ptr::NonNull::from(*node).cast(),"
+            f"AnyNodeRef::{node.name}(node) => std::ptr::NonNull::from(node.as_ref()).cast(),"
         )
     out.append("""
                 }
@@ -470,8 +642,10 @@ def write_nodekind(out: list[str], ast: Ast) -> None:
 def generate(ast: Ast) -> list[str]:
     out = []
     write_preamble(out)
+    write_ids(out, ast)
     write_owned_enum(out, ast)
     write_ref_enum(out, ast)
+    write_storage(out, ast)
     write_anynoderef(out, ast)
     write_nodekind(out, ast)
     return out

crates/ruff_python_ast/src/ast.rs

@@ -0,0 +1,92 @@
+#![allow(clippy::derive_partial_eq_without_eq)]
+
+use std::ops::{Deref, Index};
+
+use crate as ast;
+
+#[derive(Clone, Default, PartialEq)]
+pub struct Ast {
+    pub(crate) storage: ast::Storage,
+}
+
+impl Ast {
+    #[inline]
+    pub fn wrap<T>(&self, node: T) -> Node<T> {
+        Node { ast: self, node }
+    }
+
+    #[inline]
+    pub fn node<'a, I>(&'a self, id: I) -> <I as AstId>::Output<'a>
+    where
+        I: AstId,
+    {
+        id.node(self)
+    }
+}
+
+impl std::fmt::Debug for Ast {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Ast").finish()
+    }
+}
+
+pub trait AstId {
+    type Output<'a>;
+    fn node<'a>(self, ast: &'a Ast) -> Self::Output<'a>;
+}
+
+pub trait AstIdMut {
+    type Output<'a>;
+    fn node_mut<'a>(self, ast: &'a mut Ast) -> Self::Output<'a>;
+}
+
+#[derive(Clone, Copy)]
+pub struct Node<'ast, T> {
+    pub ast: &'ast Ast,
+    pub node: T,
+}
+
+impl<T> Node<'_, T> {
+    pub fn as_ref(&self) -> &T {
+        &self.node
+    }
+}
+
+impl<T> std::fmt::Debug for Node<'_, T>
+where
+    T: std::fmt::Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Node").field(&self.node).finish()
+    }
+}
+
+impl<T> Deref for Node<'_, T> {
+    type Target = T;
+    fn deref(&self) -> &Self::Target {
+        &self.node
+    }
+}
+
+impl<T> Eq for Node<'_, T> where T: Eq {}
+
+impl<T> std::hash::Hash for Node<'_, T>
+where
+    T: std::hash::Hash,
+{
+    fn hash<H>(&self, state: &mut H)
+    where
+        H: std::hash::Hasher,
+    {
+        self.node.hash(state);
+    }
+}
+
+impl<T> PartialEq for Node<'_, T>
+where
+    T: PartialEq,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.node == other.node
+    }
+}

crates/ruff_python_ast/src/comparable.rs

@@ -16,7 +16,7 @@
 //! have the same shape in that they evaluate to the same value.
 
 use crate as ast;
-use crate::{Expr, Number};
+use crate::{Expr, Node, Number};
 use std::borrow::Cow;
 use std::hash::Hash;
 
@@ -593,18 +593,14 @@ impl<'a> From<ast::LiteralExpressionRef<'a>> for ComparableLiteral<'a> {
         match literal {
             ast::LiteralExpressionRef::NoneLiteral(_) => Self::None,
             ast::LiteralExpressionRef::EllipsisLiteral(_) => Self::Ellipsis,
-            ast::LiteralExpressionRef::BooleanLiteral(ast::ExprBooleanLiteral {
-                value, ..
-            }) => Self::Bool(value),
-            ast::LiteralExpressionRef::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
-                Self::Str(value.iter().map(Into::into).collect())
+            ast::LiteralExpressionRef::BooleanLiteral(node) => Self::Bool(&node.value),
+            ast::LiteralExpressionRef::StringLiteral(node) => {
+                Self::Str(node.value.iter().map(Into::into).collect())
             }
-            ast::LiteralExpressionRef::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => {
-                Self::Bytes(value.iter().map(Into::into).collect())
-            }
-            ast::LiteralExpressionRef::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => {
-                Self::Number(value.into())
+            ast::LiteralExpressionRef::BytesLiteral(node) => {
+                Self::Bytes(node.value.iter().map(Into::into).collect())
             }
+            ast::LiteralExpressionRef::NumberLiteral(node) => Self::Number((&node.value).into()),
         }
     }
 }
@@ -1437,9 +1433,9 @@ pub enum ComparableStmt<'a> {
     Continue,
 }
 
-impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
-    fn from(stmt: &'a ast::Stmt) -> Self {
-        match stmt {
+impl<'a> From<Node<'a, ast::StmtRef<'a>>> for ComparableStmt<'a> {
+    fn from(stmt: Node<'a, ast::StmtRef<'a>>) -> Self {
+        match stmt.node {
             ast::Stmt::FunctionDef(ast::StmtFunctionDef {
                 is_async,
                 name,
@@ -1451,7 +1447,7 @@ impl<'a> From<&'a ast::Stmt> for ComparableStmt<'a> {
                 range: _,
             }) => Self::FunctionDef(StmtFunctionDef {
                 is_async: *is_async,
-                name: name.as_str(),
+                name: stmt.name().as_str(),
                 parameters: parameters.into(),
                 body: body.iter().map(Into::into).collect(),
                 decorator_list: decorator_list.iter().map(Into::into).collect(),

crates/ruff_python_ast/src/expression.rs

@@ -3,26 +3,20 @@ use std::iter::FusedIterator;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::{
-    self as ast, AnyNodeRef, AnyStringFlags, Expr, ExprBytesLiteral, ExprFString, ExprRef,
-    ExprStringLiteral, StringFlags,
+    self as ast, AnyNodeRef, AnyStringFlags, Ast, Expr, ExprBytesLiteral, ExprFString, ExprRef,
+    ExprStringLiteral, Node, StringFlags,
 };
 
-impl<'a> From<&'a Box<Expr>> for ExprRef<'a> {
-    fn from(value: &'a Box<Expr>) -> Self {
-        ExprRef::from(value.as_ref())
-    }
-}
-
 /// Unowned pendant to all the literal variants of [`ast::Expr`] that stores a
 /// reference instead of an owned value.
 #[derive(Copy, Clone, Debug, PartialEq, is_macro::Is)]
 pub enum LiteralExpressionRef<'a> {
-    StringLiteral(&'a ast::ExprStringLiteral),
-    BytesLiteral(&'a ast::ExprBytesLiteral),
-    NumberLiteral(&'a ast::ExprNumberLiteral),
-    BooleanLiteral(&'a ast::ExprBooleanLiteral),
-    NoneLiteral(&'a ast::ExprNoneLiteral),
-    EllipsisLiteral(&'a ast::ExprEllipsisLiteral),
+    StringLiteral(Node<'a, &'a ast::ExprStringLiteral>),
+    BytesLiteral(Node<'a, &'a ast::ExprBytesLiteral>),
+    NumberLiteral(Node<'a, &'a ast::ExprNumberLiteral>),
+    BooleanLiteral(Node<'a, &'a ast::ExprBooleanLiteral>),
+    NoneLiteral(Node<'a, &'a ast::ExprNoneLiteral>),
+    EllipsisLiteral(Node<'a, &'a ast::ExprEllipsisLiteral>),
 }
 
 impl Ranged for LiteralExpressionRef<'_> {
@@ -83,9 +77,9 @@ impl LiteralExpressionRef<'_> {
 /// literals, bytes literals, and f-strings.
 #[derive(Copy, Clone, Debug, PartialEq)]
 pub enum StringLike<'a> {
-    String(&'a ast::ExprStringLiteral),
-    Bytes(&'a ast::ExprBytesLiteral),
-    FString(&'a ast::ExprFString),
+    String(Node<'a, &'a ast::ExprStringLiteral>),
+    Bytes(Node<'a, &'a ast::ExprBytesLiteral>),
+    FString(Node<'a, &'a ast::ExprFString>),
 }
 
 impl<'a> StringLike<'a> {
@@ -96,18 +90,18 @@ impl<'a> StringLike<'a> {
     /// Returns an iterator over the [`StringLikePart`] contained in this string-like expression.
     pub fn parts(&self) -> StringLikePartIter<'a> {
         match self {
-            StringLike::String(expr) => StringLikePartIter::String(expr.value.iter()),
-            StringLike::Bytes(expr) => StringLikePartIter::Bytes(expr.value.iter()),
-            StringLike::FString(expr) => StringLikePartIter::FString(expr.value.iter()),
+            StringLike::String(expr) => StringLikePartIter::String(expr.ast, expr.value.iter()),
+            StringLike::Bytes(expr) => StringLikePartIter::Bytes(expr.ast, expr.value.iter()),
+            StringLike::FString(expr) => StringLikePartIter::FString(expr.ast, expr.value.iter()),
         }
     }
 
     /// Returns `true` if the string is implicitly concatenated.
     pub fn is_implicit_concatenated(self) -> bool {
         match self {
-            Self::String(ExprStringLiteral { value, .. }) => value.is_implicit_concatenated(),
-            Self::Bytes(ExprBytesLiteral { value, .. }) => value.is_implicit_concatenated(),
-            Self::FString(ExprFString { value, .. }) => value.is_implicit_concatenated(),
+            Self::String(node) => node.value.is_implicit_concatenated(),
+            Self::Bytes(node) => node.value.is_implicit_concatenated(),
+            Self::FString(node) => node.value.is_implicit_concatenated(),
         }
     }
 
@@ -120,26 +114,26 @@ impl<'a> StringLike<'a> {
     }
 }
 
-impl<'a> From<&'a ast::ExprStringLiteral> for StringLike<'a> {
-    fn from(value: &'a ast::ExprStringLiteral) -> Self {
+impl<'a> From<Node<'a, &'a ast::ExprStringLiteral>> for StringLike<'a> {
+    fn from(value: Node<'a, &'a ast::ExprStringLiteral>) -> Self {
         StringLike::String(value)
     }
 }
 
-impl<'a> From<&'a ast::ExprBytesLiteral> for StringLike<'a> {
-    fn from(value: &'a ast::ExprBytesLiteral) -> Self {
+impl<'a> From<Node<'a, &'a ast::ExprBytesLiteral>> for StringLike<'a> {
+    fn from(value: Node<'a, &'a ast::ExprBytesLiteral>) -> Self {
         StringLike::Bytes(value)
     }
 }
 
-impl<'a> From<&'a ast::ExprFString> for StringLike<'a> {
-    fn from(value: &'a ast::ExprFString) -> Self {
+impl<'a> From<Node<'a, &'a ast::ExprFString>> for StringLike<'a> {
+    fn from(value: Node<'a, &'a ast::ExprFString>) -> Self {
         StringLike::FString(value)
     }
 }
 
-impl<'a> From<&StringLike<'a>> for ExprRef<'a> {
-    fn from(value: &StringLike<'a>) -> Self {
+impl<'a> From<StringLike<'a>> for ExprRef<'a> {
+    fn from(value: StringLike<'a>) -> Self {
         match value {
             StringLike::String(expr) => ExprRef::StringLiteral(expr),
             StringLike::Bytes(expr) => ExprRef::BytesLiteral(expr),
@@ -150,12 +144,6 @@ impl<'a> From<&StringLike<'a>> for ExprRef<'a> {
 
 impl<'a> From<StringLike<'a>> for AnyNodeRef<'a> {
     fn from(value: StringLike<'a>) -> Self {
-        AnyNodeRef::from(&value)
-    }
-}
-
-impl<'a> From<&StringLike<'a>> for AnyNodeRef<'a> {
-    fn from(value: &StringLike<'a>) -> Self {
         match value {
             StringLike::String(expr) => AnyNodeRef::ExprStringLiteral(expr),
             StringLike::Bytes(expr) => AnyNodeRef::ExprBytesLiteral(expr),
@@ -164,14 +152,14 @@ impl<'a> From<&StringLike<'a>> for AnyNodeRef<'a> {
     }
 }
 
-impl<'a> TryFrom<&'a Expr> for StringLike<'a> {
+impl<'a> TryFrom<Node<'a, &'a Expr>> for StringLike<'a> {
     type Error = ();
 
-    fn try_from(value: &'a Expr) -> Result<Self, Self::Error> {
-        match value {
-            Expr::StringLiteral(value) => Ok(Self::String(value)),
-            Expr::BytesLiteral(value) => Ok(Self::Bytes(value)),
-            Expr::FString(value) => Ok(Self::FString(value)),
+    fn try_from(value: Node<'a, &'a Expr>) -> Result<Self, Self::Error> {
+        match value.node {
+            Expr::StringLiteral(v) => Ok(Self::String(value.ast.wrap(v))),
+            Expr::BytesLiteral(v) => Ok(Self::Bytes(value.ast.wrap(v))),
+            Expr::FString(v) => Ok(Self::FString(value.ast.wrap(v))),
             _ => Err(()),
         }
     }
@@ -203,9 +191,9 @@ impl Ranged for StringLike<'_> {
 /// An enum that holds a reference to an individual part of a string-like expression.
 #[derive(Copy, Clone, Debug, PartialEq)]
 pub enum StringLikePart<'a> {
-    String(&'a ast::StringLiteral),
-    Bytes(&'a ast::BytesLiteral),
-    FString(&'a ast::FString),
+    String(Node<'a, &'a ast::StringLiteral>),
+    Bytes(Node<'a, &'a ast::BytesLiteral>),
+    FString(Node<'a, &'a ast::FString>),
 }
 
 impl<'a> StringLikePart<'a> {
@@ -231,7 +219,7 @@ impl<'a> StringLikePart<'a> {
         matches!(self, Self::String(_))
     }
 
-    pub const fn as_string_literal(self) -> Option<&'a ast::StringLiteral> {
+    pub const fn as_string_literal(self) -> Option<Node<'a, &'a ast::StringLiteral>> {
         match self {
             StringLikePart::String(value) => Some(value),
             _ => None,
@@ -243,30 +231,24 @@ impl<'a> StringLikePart<'a> {
     }
 }
 
-impl<'a> From<&'a ast::StringLiteral> for StringLikePart<'a> {
-    fn from(value: &'a ast::StringLiteral) -> Self {
+impl<'a> From<Node<'a, &'a ast::StringLiteral>> for StringLikePart<'a> {
+    fn from(value: Node<'a, &'a ast::StringLiteral>) -> Self {
         StringLikePart::String(value)
     }
 }
 
-impl<'a> From<&'a ast::BytesLiteral> for StringLikePart<'a> {
-    fn from(value: &'a ast::BytesLiteral) -> Self {
+impl<'a> From<Node<'a, &'a ast::BytesLiteral>> for StringLikePart<'a> {
+    fn from(value: Node<'a, &'a ast::BytesLiteral>) -> Self {
         StringLikePart::Bytes(value)
     }
 }
 
-impl<'a> From<&'a ast::FString> for StringLikePart<'a> {
-    fn from(value: &'a ast::FString) -> Self {
+impl<'a> From<Node<'a, &'a ast::FString>> for StringLikePart<'a> {
+    fn from(value: Node<'a, &'a ast::FString>) -> Self {
         StringLikePart::FString(value)
     }
 }
 
-impl<'a> From<&StringLikePart<'a>> for AnyNodeRef<'a> {
-    fn from(value: &StringLikePart<'a>) -> Self {
-        AnyNodeRef::from(*value)
-    }
-}
-
 impl<'a> From<StringLikePart<'a>> for AnyNodeRef<'a> {
     fn from(value: StringLikePart<'a>) -> Self {
         match value {
@@ -292,9 +274,9 @@ impl Ranged for StringLikePart<'_> {
 /// This is created by the [`StringLike::parts`] method.
 #[derive(Clone)]
 pub enum StringLikePartIter<'a> {
-    String(std::slice::Iter<'a, ast::StringLiteral>),
-    Bytes(std::slice::Iter<'a, ast::BytesLiteral>),
-    FString(std::slice::Iter<'a, ast::FStringPart>),
+    String(&'a Ast, std::slice::Iter<'a, ast::StringLiteral>),
+    Bytes(&'a Ast, std::slice::Iter<'a, ast::BytesLiteral>),
+    FString(&'a Ast, std::slice::Iter<'a, ast::FStringPart>),
 }
 
 impl<'a> Iterator for StringLikePartIter<'a> {
@@ -302,15 +284,19 @@ impl<'a> Iterator for StringLikePartIter<'a> {
 
     fn next(&mut self) -> Option<Self::Item> {
         let part = match self {
-            StringLikePartIter::String(inner) => StringLikePart::String(inner.next()?),
-            StringLikePartIter::Bytes(inner) => StringLikePart::Bytes(inner.next()?),
-            StringLikePartIter::FString(inner) => {
+            StringLikePartIter::String(ast, inner) => {
+                StringLikePart::String(ast.wrap(inner.next()?))
+            }
+            StringLikePartIter::Bytes(ast, inner) => StringLikePart::Bytes(ast.wrap(inner.next()?)),
+            StringLikePartIter::FString(ast, inner) => {
                 let part = inner.next()?;
                 match part {
                     ast::FStringPart::Literal(string_literal) => {
-                        StringLikePart::String(string_literal)
+                        StringLikePart::String(ast.wrap(string_literal))
+                    }
+                    ast::FStringPart::FString(f_string) => {
+                        StringLikePart::FString(ast.wrap(f_string))
                     }
-                    ast::FStringPart::FString(f_string) => StringLikePart::FString(f_string),
                 }
             }
         };
@@ -320,9 +306,9 @@ impl<'a> Iterator for StringLikePartIter<'a> {
 
     fn size_hint(&self) -> (usize, Option<usize>) {
         match self {
-            StringLikePartIter::String(inner) => inner.size_hint(),
-            StringLikePartIter::Bytes(inner) => inner.size_hint(),
-            StringLikePartIter::FString(inner) => inner.size_hint(),
+            StringLikePartIter::String(_, inner) => inner.size_hint(),
+            StringLikePartIter::Bytes(_, inner) => inner.size_hint(),
+            StringLikePartIter::FString(_, inner) => inner.size_hint(),
         }
     }
 }
@@ -330,15 +316,21 @@ impl<'a> Iterator for StringLikePartIter<'a> {
 impl DoubleEndedIterator for StringLikePartIter<'_> {
     fn next_back(&mut self) -> Option<Self::Item> {
         let part = match self {
-            StringLikePartIter::String(inner) => StringLikePart::String(inner.next_back()?),
-            StringLikePartIter::Bytes(inner) => StringLikePart::Bytes(inner.next_back()?),
-            StringLikePartIter::FString(inner) => {
+            StringLikePartIter::String(ast, inner) => {
+                StringLikePart::String(ast.wrap(inner.next_back()?))
+            }
+            StringLikePartIter::Bytes(ast, inner) => {
+                StringLikePart::Bytes(ast.wrap(inner.next_back()?))
+            }
+            StringLikePartIter::FString(ast, inner) => {
                 let part = inner.next_back()?;
                 match part {
                     ast::FStringPart::Literal(string_literal) => {
-                        StringLikePart::String(string_literal)
+                        StringLikePart::String(ast.wrap(string_literal))
+                    }
+                    ast::FStringPart::FString(f_string) => {
+                        StringLikePart::FString(ast.wrap(f_string))
                     }
-                    ast::FStringPart::FString(f_string) => StringLikePart::FString(f_string),
                 }
             }
         };

crates/ruff_python_ast/src/lib.rs

@@ -1,11 +1,13 @@
 use std::ffi::OsStr;
 use std::path::Path;
 
+pub use ast::{Ast, Node};
 pub use expression::*;
 pub use generated::*;
 pub use int::*;
 pub use nodes::*;
 
+pub mod ast;
 pub mod comparable;
 pub mod docstrings;
 mod expression;

crates/ruff_python_ast/src/nodes.rs

@@ -13,28 +13,90 @@ use itertools::Itertools;
 
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
+use crate::ast::AstId;
 use crate::name::Name;
 use crate::{
     int,
     str::Quote,
     str_prefix::{AnyStringPrefix, ByteStringPrefix, FStringPrefix, StringLiteralPrefix},
-    ExceptHandler, Expr, FStringElement, LiteralExpressionRef, Pattern, Stmt, TypeParam,
+    Ast, DecoratorId, ExceptHandler, Expr, ExprId, FStringElement, IdentifierId,
+    LiteralExpressionRef, Node, ParametersId, Pattern, Stmt, StmtId, TypeParam, TypeParamsId,
 };
 
+macro_rules! accessor {
+    ($ty:ty, $field:ident, $ref_ty:ty) => {
+        impl $ty {
+            pub fn $field<'a>(&'a self, ast: &'a Ast) -> Node<'a, <$ref_ty as AstId>::Output<'a>> {
+                ast.wrap(ast.node(self.$field))
+            }
+        }
+
+        impl<'a> Node<'a, &'a $ty> {
+            pub fn $field(self) -> Node<'a, <$ref_ty as AstId>::Output<'a>> {
+                self.node.$field(self.ast)
+            }
+        }
+    };
+}
+
+macro_rules! option_accessor {
+    ($ty:ty, $field:ident, $ref_ty:ty) => {
+        impl $ty {
+            pub fn $field<'a>(
+                &'a self,
+                ast: &'a Ast,
+            ) -> Option<Node<'a, <$ref_ty as AstId>::Output<'a>>> {
+                self.$field.map(|id| ast.wrap(ast.node(*id)))
+            }
+        }
+
+        impl<'a> Node<'a, &'a $ty> {
+            pub fn $field(self) -> Option<Node<'a, <$ref_ty as AstId>::Output<'a>>> {
+                self.node.$field(self.ast)
+            }
+        }
+    };
+}
+
+macro_rules! vec_accessor {
+    ($ty:ty, $field:ident, $ref_ty:ty) => {
+        impl $ty {
+            pub fn $field<'a>(
+                &'a self,
+                ast: &'a Ast,
+            ) -> impl Iterator<Item = Node<'a, <$ref_ty as AstId>::Output<'a>>> + 'a {
+                self.$field.iter().map(|id| ast.wrap(ast.node(*id)))
+            }
+        }
+
+        impl<'a> Node<'a, &'a $ty> {
+            pub fn $field(
+                self,
+            ) -> impl Iterator<Item = Node<'a, <$ref_ty as AstId>::Output<'a>>> + 'a {
+                self.node.$field(self.ast)
+            }
+        }
+    };
+}
+
 /// See also [Module](https://docs.python.org/3/library/ast.html#ast.Module)
 #[derive(Clone, Debug, PartialEq)]
 pub struct ModModule {
     pub range: TextRange,
-    pub body: Vec<Stmt>,
+    pub body: Vec<StmtId>,
 }
 
+vec_accessor!(ModModule, body, StmtId);
+
 /// See also [Expression](https://docs.python.org/3/library/ast.html#ast.Expression)
 #[derive(Clone, Debug, PartialEq)]
 pub struct ModExpression {
     pub range: TextRange,
-    pub body: Box<Expr>,
+    pub body: ExprId,
 }
 
+accessor!(ModExpression, body, ExprId);
+
 /// An AST node used to represent a IPython escape command at the statement level.
 ///
 /// For example,
@@ -104,20 +166,27 @@ pub struct StmtIpyEscapeCommand {
 pub struct StmtFunctionDef {
     pub range: TextRange,
     pub is_async: bool,
-    pub decorator_list: Vec<Decorator>,
-    pub name: Identifier,
-    pub type_params: Option<Box<TypeParams>>,
-    pub parameters: Box<Parameters>,
-    pub returns: Option<Box<Expr>>,
-    pub body: Vec<Stmt>,
+    pub decorator_list: Vec<DecoratorId>,
+    pub name: IdentifierId,
+    pub type_params: Option<TypeParamsId>,
+    pub parameters: ParametersId,
+    pub returns: Option<ExprId>,
+    pub body: Vec<StmtId>,
 }
 
+vec_accessor!(StmtFunctionDef, decorator_list, DecoratorId);
+accessor!(StmtFunctionDef, name, IdentifierId);
+option_accessor!(StmtFunctionDef, type_params, TypeParamsId);
+accessor!(StmtFunctionDef, parameters, ParametersId);
+option_accessor!(StmtFunctionDef, returns, ExprId);
+vec_accessor!(StmtFunctionDef, body, StmtId);
+
 /// See also [ClassDef](https://docs.python.org/3/library/ast.html#ast.ClassDef)
 #[derive(Clone, Debug, PartialEq)]
 pub struct StmtClassDef {
     pub range: TextRange,
     pub decorator_list: Vec<Decorator>,
-    pub name: Identifier,
+    pub name: IdentifierId,
     pub type_params: Option<Box<TypeParams>>,
     pub arguments: Option<Box<Arguments>>,
     pub body: Vec<Stmt>,

crates/ruff_python_ast/src/visitor/source_order.rs

@@ -1,9 +1,9 @@
 use crate::AnyNodeRef;
 use crate::{
     Alias, Arguments, BoolOp, BytesLiteral, CmpOp, Comprehension, Decorator, ElifElseClause,
-    ExceptHandler, Expr, FString, FStringElement, Keyword, MatchCase, Mod, Operator, Parameter,
-    ParameterWithDefault, Parameters, Pattern, PatternArguments, PatternKeyword, Singleton, Stmt,
-    StringLiteral, TypeParam, TypeParams, UnaryOp, WithItem,
+    ExceptHandler, Expr, FString, FStringElement, Keyword, MatchCase, Mod, ModId, Node, Operator,
+    Parameter, ParameterWithDefault, Parameters, Pattern, PatternArguments, PatternKeyword,
+    Singleton, Stmt, StringLiteral, TypeParam, TypeParams, UnaryOp, WithItem,
 };
 
 /// Visitor that traverses all nodes recursively in the order they appear in the source.
@@ -20,8 +20,8 @@ pub trait SourceOrderVisitor<'a> {
     fn leave_node(&mut self, _node: AnyNodeRef<'a>) {}
 
     #[inline]
-    fn visit_mod(&mut self, module: &'a Mod) {
-        walk_module(self, module);
+    fn visit_mod(&mut self, module: Node<'a, ModId>) {
+        walk_module(self, module.node());
     }
 
     #[inline]
@@ -172,7 +172,7 @@ pub trait SourceOrderVisitor<'a> {
     }
 }
 
-pub fn walk_module<'a, V>(visitor: &mut V, module: &'a Mod)
+pub fn walk_module<'a, V>(visitor: &mut V, module: Node<'a, &'a Mod>)
 where
     V: SourceOrderVisitor<'a> + ?Sized,
 {

crates/ruff_python_formatter/generate.py

@@ -15,6 +15,13 @@ def rustfmt(code: str) -> str:
     return check_output(["rustfmt", "--emit=stdout"], input=code, text=True)
 
 
+# Which nodes have been migrated over to the new IndexVec representation?
+indexed_nodes = {
+    "ModModule",
+    "ModExpression",
+}
+
+
 # %%
 # Read nodes
 
@@ -130,14 +137,18 @@ use ruff_python_ast as ast;
 
 """
 for node in nodes:
+    if node in indexed_nodes:
+        node_type = f"ast::Node<&ast::{node}>"
+    else:
+        node_type = f"ast::{node}"
     text = f"""
-        impl FormatRule<ast::{node}, PyFormatContext<'_>>
+        impl FormatRule<{node_type}, PyFormatContext<'_>>
             for crate::{groups[group_for_node(node)]}::{to_camel_case(node)}::Format{node}
         {{
             #[inline]
             fn fmt(
                 &self,
-                node: &ast::{node},
+                node: &{node_type},
                 f: &mut PyFormatter,
             ) -> FormatResult<()> {{
                 FormatNodeRule::<ast::{node}>::fmt(self, node, f)

crates/ruff_python_formatter/src/cli.rs

@@ -71,7 +71,8 @@ pub fn format_and_debug_print(source: &str, cli: &Cli, source_path: &Path) -> Re
     }
     if cli.print_comments {
         // Print preceding, following and enclosing nodes
-        let decorated_comments = collect_comments(parsed.syntax(), source_code, &comment_ranges);
+        let decorated_comments =
+            collect_comments(parsed.syntax().node(), source_code, &comment_ranges);
         if !decorated_comments.is_empty() {
             println!("# Comment decoration: Range, Preceding, Following, Enclosing, Comment");
         }

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,7 +1,8 @@
 use ast::helpers::comment_indentation_after;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{
-    self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Parameter, Parameters, StringLike,
+    self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Node, Parameter, Parameters,
+    StringLike,
 };
 use ruff_python_trivia::{
     find_only_token_in_range, first_non_trivia_token, indentation_at_offset, BackwardsTokenizer,
@@ -966,7 +967,7 @@ fn handle_trailing_binary_like_comment<'a>(
 ///
 /// Comments of an all empty module are leading module comments
 fn handle_trailing_module_comment<'a>(
-    module: &'a ModModule,
+    module: Node<'a, &'a ModModule>,
     comment: DecoratedComment<'a>,
 ) -> CommentPlacement<'a> {
     if comment.preceding_node().is_none() && comment.following_node().is_none() {

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -18,7 +18,7 @@ use crate::comments::{CommentsMap, SourceComment};
 /// Collect the preceding, following and enclosing node for each comment without applying
 /// [`place_comment`] for debugging.
 pub(crate) fn collect_comments<'a>(
-    root: &'a Mod,
+    root: Mod,
     source_code: SourceCode<'a>,
     comment_ranges: &'a CommentRanges,
 ) -> Vec<DecoratedComment<'a>> {

crates/ruff_python_formatter/src/lib.rs

@@ -4,7 +4,7 @@ use tracing::Level;
 pub use range::format_range;
 use ruff_formatter::prelude::*;
 use ruff_formatter::{format, write, FormatError, Formatted, PrintError, Printed, SourceCode};
-use ruff_python_ast::{AnyNodeRef, Mod};
+use ruff_python_ast::{AnyNodeRef, Mod, ModId};
 use ruff_python_parser::{parse, AsMode, ParseError, Parsed};
 use ruff_python_trivia::CommentRanges;
 use ruff_text_size::Ranged;
@@ -119,7 +119,7 @@ pub fn format_module_source(
 }
 
 pub fn format_module_ast<'a>(
-    parsed: &'a Parsed<Mod>,
+    parsed: &'a Parsed<ModId>,
     comment_ranges: &'a CommentRanges,
     source: &'a str,
     options: PyFormatOptions,

crates/ruff_python_formatter/src/module/mod.rs

@@ -1,4 +1,4 @@
-use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule};
+use ruff_formatter::FormatOwnedWithRule;
 use ruff_python_ast::Mod;
 
 use crate::prelude::*;
@@ -9,8 +9,8 @@ pub(crate) mod mod_module;
 #[derive(Default)]
 pub struct FormatMod;
 
-impl FormatRule<Mod, PyFormatContext<'_>> for FormatMod {
-    fn fmt(&self, item: &Mod, f: &mut PyFormatter) -> FormatResult<()> {
+impl FormatRule<Mod<'_>, PyFormatContext<'_>> for FormatMod {
+    fn fmt(&self, item: &Mod<'_>, f: &mut PyFormatter) -> FormatResult<()> {
         match item {
             Mod::Module(x) => x.format().fmt(f),
             Mod::Expression(x) => x.format().fmt(f),
@@ -18,16 +18,8 @@ impl FormatRule<Mod, PyFormatContext<'_>> for FormatMod {
     }
 }
 
-impl<'ast> AsFormat<PyFormatContext<'ast>> for Mod {
-    type Format<'a> = FormatRefWithRule<'a, Mod, FormatMod, PyFormatContext<'ast>>;
-
-    fn format(&self) -> Self::Format<'_> {
-        FormatRefWithRule::new(self, FormatMod)
-    }
-}
-
-impl<'ast> IntoFormat<PyFormatContext<'ast>> for Mod {
-    type Format = FormatOwnedWithRule<Mod, FormatMod, PyFormatContext<'ast>>;
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for Mod<'ast> {
+    type Format = FormatOwnedWithRule<Mod<'ast>, FormatMod, PyFormatContext<'ast>>;
 
     fn into_format(self) -> Self::Format {
         FormatOwnedWithRule::new(self, FormatMod)

crates/ruff_python_formatter/src/module/mod_expression.rs

@@ -1,4 +1,4 @@
-use ruff_python_ast::ModExpression;
+use ruff_python_ast::{ModExpression, Node};
 
 use crate::prelude::*;
 

crates/ruff_python_formatter/src/module/mod_module.rs

@@ -1,5 +1,5 @@
 use ruff_formatter::write;
-use ruff_python_ast::ModModule;
+use ruff_python_ast::{ModModule, Node};
 use ruff_python_trivia::lines_after;
 
 use crate::prelude::*;

crates/ruff_python_parser/src/lib.rs

@@ -73,7 +73,8 @@ pub use crate::token::{Token, TokenKind};
 use crate::parser::Parser;
 
 use ruff_python_ast::{
-    Expr, Mod, ModExpression, ModModule, PySourceType, StringFlags, StringLiteral, Suite,
+    Ast, Expr, Mod, ModExpression, ModId, ModModule, ModModuleId, PySourceType, StringFlags,
+    StringLiteral, Suite,
 };
 use ruff_python_trivia::CommentRanges;
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -109,7 +110,7 @@ pub mod typing;
 /// let module = parse_module(source);
 /// assert!(module.is_ok());
 /// ```
-pub fn parse_module(source: &str) -> Result<Parsed<ModModule>, ParseError> {
+pub fn parse_module(source: &str) -> Result<Parsed<ModModuleId>, ParseError> {
     Parser::new(source, Mode::Module)
         .parse()
         .try_into_module()
@@ -132,7 +133,7 @@ pub fn parse_module(source: &str) -> Result<Parsed<ModModule>, ParseError> {
 /// let expr = parse_expression("1 + 2");
 /// assert!(expr.is_ok());
 /// ```
-pub fn parse_expression(source: &str) -> Result<Parsed<ModExpression>, ParseError> {
+pub fn parse_expression(source: &str) -> Result<Parsed<ModExpressionId>, ParseError> {
     Parser::new(source, Mode::Expression)
         .parse()
         .try_into_expression()
@@ -159,7 +160,7 @@ pub fn parse_expression(source: &str) -> Result<Parsed<ModExpression>, ParseErro
 pub fn parse_expression_range(
     source: &str,
     range: TextRange,
-) -> Result<Parsed<ModExpression>, ParseError> {
+) -> Result<Parsed<ModExpressionId>, ParseError> {
     let source = &source[..range.end().to_usize()];
     Parser::new_starts_at(source, Mode::Expression, range.start())
         .parse()
@@ -273,7 +274,7 @@ pub fn parse_string_annotation(
 /// let parsed = parse(source, Mode::Ipython);
 /// assert!(parsed.is_ok());
 /// ```
-pub fn parse(source: &str, mode: Mode) -> Result<Parsed<Mod>, ParseError> {
+pub fn parse(source: &str, mode: Mode) -> Result<Parsed<ModId>, ParseError> {
     parse_unchecked(source, mode).into_result()
 }
 
@@ -281,12 +282,12 @@ pub fn parse(source: &str, mode: Mode) -> Result<Parsed<Mod>, ParseError> {
 ///
 /// This is same as the [`parse`] function except that it doesn't check for any [`ParseError`]
 /// and returns the [`Parsed`] as is.
-pub fn parse_unchecked(source: &str, mode: Mode) -> Parsed<Mod> {
+pub fn parse_unchecked(source: &str, mode: Mode) -> Parsed<ModId> {
     Parser::new(source, mode).parse()
 }
 
 /// Parse the given Python source code using the specified [`PySourceType`].
-pub fn parse_unchecked_source(source: &str, source_type: PySourceType) -> Parsed<ModModule> {
+pub fn parse_unchecked_source(source: &str, source_type: PySourceType) -> Parsed<ModModuleId> {
     // SAFETY: Safe because `PySourceType` always parses to a `ModModule`
     Parser::new(source, source_type.as_mode())
         .parse()
@@ -297,15 +298,26 @@ pub fn parse_unchecked_source(source: &str, source_type: PySourceType) -> Parsed
 /// Represents the parsed source code.
 #[derive(Debug, PartialEq, Clone)]
 pub struct Parsed<T> {
+    ast: Ast,
     syntax: T,
     tokens: Tokens,
     errors: Vec<ParseError>,
 }
 
-impl<T> Parsed<T> {
+impl<T> Parsed<T>
+where
+    T: Clone,
+{
     /// Returns the syntax node represented by this parsed output.
-    pub fn syntax(&self) -> &T {
-        &self.syntax
+    pub fn syntax(&self) -> Node<T> {
+        self.ast.wrap(self.syntax.clone())
+    }
+}
+
+impl<T> Parsed<T> {
+    /// Returns the AST for the parsed output.
+    pub fn ast(&self) -> &Ast {
+        &self.ast
     }
 
     /// Returns all the tokens for the parsed output.
@@ -318,11 +330,6 @@ impl<T> Parsed<T> {
         &self.errors
     }
 
-    /// Consumes the [`Parsed`] output and returns the contained syntax node.
-    pub fn into_syntax(self) -> T {
-        self.syntax
-    }
-
     /// Consumes the [`Parsed`] output and returns a list of syntax errors found during parsing.
     pub fn into_errors(self) -> Vec<ParseError> {
         self.errors
@@ -354,7 +361,7 @@ impl<T> Parsed<T> {
     }
 }
 
-impl Parsed<Mod> {
+impl Parsed<ModId> {
     /// Attempts to convert the [`Parsed<Mod>`] into a [`Parsed<ModModule>`].
     ///
     /// This method checks if the `syntax` field of the output is a [`Mod::Module`]. If it is, the
@@ -362,14 +369,15 @@ impl Parsed<Mod> {
     /// returns [`None`].
     ///
     /// [`Some(Parsed<ModModule>)`]: Some
-    pub fn try_into_module(self) -> Option<Parsed<ModModule>> {
+    pub fn try_into_module(self) -> Option<Parsed<ModModuleId>> {
         match self.syntax {
-            Mod::Module(module) => Some(Parsed {
+            ModId::Module(module) => Some(Parsed {
+                ast: self.ast,
                 syntax: module,
                 tokens: self.tokens,
                 errors: self.errors,
             }),
-            Mod::Expression(_) => None,
+            ModId::Expression(_) => None,
         }
     }
 
@@ -380,10 +388,11 @@ impl Parsed<Mod> {
     /// Otherwise, it returns [`None`].
     ///
     /// [`Some(Parsed<ModExpression>)`]: Some
-    pub fn try_into_expression(self) -> Option<Parsed<ModExpression>> {
+    pub fn try_into_expression(self) -> Option<Parsed<ModExpressionId>> {
         match self.syntax {
-            Mod::Module(_) => None,
-            Mod::Expression(expression) => Some(Parsed {
+            ModId::Module(_) => None,
+            ModId::Expression(expression) => Some(Parsed {
+                ast: self.ast,
                 syntax: expression,
                 tokens: self.tokens,
                 errors: self.errors,
@@ -392,32 +401,22 @@ impl Parsed<Mod> {
     }
 }
 
-impl Parsed<ModModule> {
+impl Parsed<ModModuleId> {
     /// Returns the module body contained in this parsed output as a [`Suite`].
     pub fn suite(&self) -> &Suite {
-        &self.syntax.body
-    }
-
-    /// Consumes the [`Parsed`] output and returns the module body as a [`Suite`].
-    pub fn into_suite(self) -> Suite {
-        self.syntax.body
+        &self.ast[self.syntax].body
     }
 }
 
-impl Parsed<ModExpression> {
+impl Parsed<ModExpressionId> {
     /// Returns the expression contained in this parsed output.
     pub fn expr(&self) -> &Expr {
-        &self.syntax.body
+        &self.ast[self.syntax].body
     }
 
     /// Returns a mutable reference to the expression contained in this parsed output.
     pub fn expr_mut(&mut self) -> &mut Expr {
-        &mut self.syntax.body
-    }
-
-    /// Consumes the [`Parsed`] output and returns the contained [`Expr`].
-    pub fn into_expr(self) -> Expr {
-        *self.syntax.body
+        &mut self.ast[self.syntax].body
     }
 }
 

crates/ruff_python_parser/src/parser/mod.rs

@@ -2,7 +2,7 @@ use std::cmp::Ordering;
 
 use bitflags::bitflags;
 
-use ruff_python_ast::{Mod, ModExpression, ModModule};
+use ruff_python_ast::{Ast, ModExpression, ModId, ModModule};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::parser::expression::ExpressionContext;
@@ -47,6 +47,9 @@ pub(crate) struct Parser<'src> {
 
     /// The start offset in the source code from which to start parsing at.
     start_offset: TextSize,
+
+    /// The AST that we are constructing.
+    ast: Ast,
 }
 
 impl<'src> Parser<'src> {
@@ -68,16 +71,21 @@ impl<'src> Parser<'src> {
             prev_token_end: TextSize::new(0),
             start_offset,
             current_token_id: TokenId::default(),
+            ast: Ast::default(),
         }
     }
 
     /// Consumes the [`Parser`] and returns the parsed [`Parsed`].
-    pub(crate) fn parse(mut self) -> Parsed<Mod> {
+    pub(crate) fn parse(mut self) -> Parsed<ModId> {
         let syntax = match self.mode {
             Mode::Expression | Mode::ParenthesizedExpression => {
-                Mod::Expression(self.parse_single_expression())
+                let expr = self.parse_single_expression();
+                self.ast.add_mod_expression(expr)
+            }
+            Mode::Module | Mode::Ipython => {
+                let module = self.parse_module();
+                self.ast.add_mod_module(module)
             }
-            Mode::Module | Mode::Ipython => Mod::Module(self.parse_module()),
         };
 
         self.finish(syntax)
@@ -140,7 +148,7 @@ impl<'src> Parser<'src> {
         }
     }
 
-    fn finish(self, syntax: Mod) -> Parsed<Mod> {
+    fn finish(self, syntax: ModId) -> Parsed<ModId> {
         assert_eq!(
             self.current_token_kind(),
             TokenKind::EndOfFile,
@@ -156,6 +164,7 @@ impl<'src> Parser<'src> {
         // always results in a parse error.
         if lex_errors.is_empty() {
             return Parsed {
+                ast: self.ast,
                 syntax,
                 tokens: Tokens::new(tokens),
                 errors: parse_errors,
@@ -187,6 +196,7 @@ impl<'src> Parser<'src> {
         merged.extend(lex_errors.map(ParseError::from));
 
         Parsed {
+            ast: self.ast,
             syntax,
             tokens: Tokens::new(tokens),
             errors: merged,

crates/ruff_python_parser/src/typing.rs

@@ -2,7 +2,7 @@
 
 use ruff_python_ast::relocate::relocate_expr;
 use ruff_python_ast::str::raw_contents;
-use ruff_python_ast::{Expr, ExprStringLiteral, ModExpression, StringLiteral};
+use ruff_python_ast::{Expr, ExprStringLiteral, ModExpressionId, StringLiteral};
 use ruff_text_size::Ranged;
 
 use crate::{parse_expression, parse_string_annotation, ParseError, Parsed};
@@ -11,12 +11,12 @@ type AnnotationParseResult = Result<ParsedAnnotation, ParseError>;
 
 #[derive(Debug)]
 pub struct ParsedAnnotation {
-    parsed: Parsed<ModExpression>,
+    parsed: Parsed<ModExpressionId>,
     kind: AnnotationKind,
 }
 
 impl ParsedAnnotation {
-    pub fn parsed(&self) -> &Parsed<ModExpression> {
+    pub fn parsed(&self) -> &Parsed<ModExpressionId> {
         &self.parsed
     }
 

crates/ruff_python_semantic/src/imports.rs

@@ -205,7 +205,6 @@ impl serde::Serialize for NameImport {
 impl<'de> serde::de::Deserialize<'de> for NameImports {
     fn deserialize<D: serde::de::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
         use ruff_python_ast::{self as ast, Stmt};
-        use ruff_python_parser::Parsed;
 
         struct AnyNameImportsVisitor;
 
@@ -217,10 +216,8 @@ impl<'de> serde::de::Deserialize<'de> for NameImports {
             }
 
             fn visit_str<E: serde::de::Error>(self, value: &str) -> Result<Self::Value, E> {
-                let body = ruff_python_parser::parse_module(value)
-                    .map(Parsed::into_suite)
-                    .map_err(E::custom)?;
-                let [stmt] = body.as_slice() else {
+                let body = ruff_python_parser::parse_module(value).map_err(E::custom)?;
+                let [stmt] = body.suite().as_slice() else {
                     return Err(E::custom("Expected a single statement"));
                 };