more robust AST node finding

Author: carljm

crates/ty_python_semantic/resources/mdtest/dataclasses/dataclasses.md

@@ -126,6 +126,15 @@ class E:
     x: int = field(default=3)
     # error: [dataclass-field-order]
     y: str
+
+import sys
+
+@dataclass
+class F:
+    x: int = 1
+    if sys.version_info > (3, 7):
+        # error: [dataclass-field-order]
+        y: str
 ```
 
 Pure class attributes (`ClassVar`) are not included in the signature of `__init__`:

crates/ty_python_semantic/src/semantic_index/definition.rs

@@ -709,6 +709,13 @@ impl DefinitionKind<'_> {
         }
     }
 
+    pub(crate) fn as_annotated_assignment(&self) -> Option<&AnnotatedAssignmentDefinitionKind> {
+        match self {
+            DefinitionKind::AnnotatedAssignment(annassign) => Some(annassign),
+            _ => None,
+        }
+    }
+
     /// Returns the [`TextRange`] of the definition target.
     ///
     /// A definition target would mainly be the node representing the place being defined i.e.,

crates/ty_python_semantic/src/types/infer.rs

@@ -1416,24 +1416,34 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     }
                 }
 
-                // Report field ordering violations
-                for name in required_after_default_field_names {
-                    if let Some(field_node) = class_node.body.iter().find_map(|stmt| {
-                        if let ast::Stmt::AnnAssign(ann_assign) = stmt {
-                            if let ast::Expr::Name(name_expr) = ann_assign.target.as_ref() {
-                                if name_expr.id == *name {
-                                    return Some(ann_assign.target.as_ref());
-                                }
-                            }
-                        }
-                        None
-                    }) {
-                        if let Some(builder) =
-                            self.context.report_lint(&DATACLASS_FIELD_ORDER, field_node)
+                if !required_after_default_field_names.is_empty() {
+                    // Report field ordering violations
+                    let body_scope = class.body_scope(self.db()).file_scope_id(self.db());
+                    let use_def_map = self.index.use_def_map(body_scope);
+                    let place_table = self.index.place_table(body_scope);
+
+                    for name in required_after_default_field_names {
+                        let Some(symbol_id) = place_table.symbol_id(name.as_str()) else {
+                            continue;
+                        };
+                        if let Some(ann_assign) = use_def_map
+                            .end_of_scope_symbol_declarations(symbol_id)
+                            .find_map(|decl_with_constraints| {
+                                decl_with_constraints
+                                    .declaration
+                                    .definition()?
+                                    .kind(self.db())
+                                    .as_annotated_assignment()
+                            })
                         {
-                            builder.into_diagnostic(format_args!(
+                            if let Some(builder) = self.context.report_lint(
+                                &DATACLASS_FIELD_ORDER,
+                                ann_assign.target(self.module()),
+                            ) {
+                                builder.into_diagnostic(format_args!(
                                 "Required field `{name}` cannot be defined after fields with default values",
                             ));
+                            }
                         }
                     }
                 }