impl

Author: asukaminato0721

pyrefly/lib/alt/class/django.rs

@@ -31,6 +31,9 @@ use crate::alt::class::class_field::ClassField;
 use crate::alt::class::enums::VALUE_PROP;
 use crate::alt::types::class_metadata::ClassSynthesizedField;
 use crate::alt::types::class_metadata::ClassSynthesizedFields;
+use crate::binding::binding::ClassFieldDefinition;
+use crate::binding::binding::ExprOrBinding;
+use crate::binding::binding::KeyClassField;
 use crate::binding::binding::KeyExport;
 use crate::types::simplify::unions;
 
@@ -48,6 +51,16 @@ const NULL: Name = Name::new_static("null");
 const MANY_TO_MANY_FIELD: Name = Name::new_static("ManyToManyField");
 const MODEL: Name = Name::new_static("Model");
 const MANYRELATEDMANAGER: Name = Name::new_static("ManyRelatedManager");
+const ONE_TO_ONE_FIELD: Name = Name::new_static("OneToOneField");
+const RELATED_NAME: Name = Name::new_static("related_name");
+const RELATED_MANAGER: Name = Name::new_static("RelatedManager");
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum DjangoRelationKind {
+    ForeignKey,
+    OneToOne,
+    ManyToMany,
+}
 
 impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
     pub fn get_django_field_type(
@@ -184,6 +197,36 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         Some(manager_type)
     }
 
+    // Get RelatedManager class from django stubs
+    fn get_related_manager_type(&self, target_model_type: Type) -> Option<Type> {
+        let django_related_module = ModuleName::django_models_fields_related_descriptors();
+        let django_related_module_exports = self.exports.get(django_related_module).finding()?;
+        if !django_related_module_exports
+            .exports(self.exports)
+            .contains_key(&RELATED_MANAGER)
+        {
+            return None;
+        }
+
+        let manager_class_type =
+            self.get_from_export(django_related_module, None, &KeyExport(RELATED_MANAGER));
+
+        let manager_class = match manager_class_type.as_ref() {
+            Type::ClassDef(cls) => cls,
+            _ => return None,
+        };
+
+        let targs_vec = vec![target_model_type];
+        let manager_type = self.specialize(
+            manager_class,
+            targs_vec,
+            TextRange::default(),
+            &self.error_swallower(),
+        );
+
+        Some(manager_type)
+    }
+
     fn resolve_target(&self, to_expr: &Expr, class: &Class) -> Option<Type> {
         match to_expr {
             // Use expr_infer to resolve the name in the current scope
@@ -242,6 +285,13 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         )
     }
 
+    pub fn is_one_to_one_field(&self, field: &Class) -> bool {
+        field.has_toplevel_qname(
+            ModuleName::django_models_fields_related().as_str(),
+            ONE_TO_ONE_FIELD.as_str(),
+        )
+    }
+
     pub fn get_django_enum_synthesized_fields(
         &self,
         cls: &Class,
@@ -481,6 +531,185 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             }
         }
 
+        if let Some(reverse_fields) = self.get_django_reverse_relationship_synthesized_fields(cls) {
+            for (name, field) in reverse_fields.fields() {
+                fields.insert(name.clone(), field.clone());
+            }
+        }
+
         Some(ClassSynthesizedFields::new(fields))
     }
+
+    fn get_django_reverse_relationship_synthesized_fields(
+        &self,
+        cls: &Class,
+    ) -> Option<ClassSynthesizedFields> {
+        let mut fields = SmallMap::new();
+
+        for field_idx in self.bindings().keys::<KeyClassField>() {
+            let binding = self.bindings().get(field_idx);
+            let Some(source_class) = &self.get_idx(binding.class_idx).0 else {
+                continue;
+            };
+            if !self.get_metadata_for_class(source_class).is_django_model() {
+                continue;
+            }
+
+            let expr = match &binding.definition {
+                ClassFieldDefinition::AssignedInBody {
+                    value: ExprOrBinding::Expr(expr),
+                    ..
+                } => expr,
+                _ => continue,
+            };
+            let call_expr = match expr.as_call_expr() {
+                Some(call_expr) => call_expr,
+                None => continue,
+            };
+
+            let relation_kind = match self.django_relation_kind(expr) {
+                Some(kind) => kind,
+                None => continue,
+            };
+
+            let Some(to_expr) = call_expr.arguments.args.first() else {
+                continue;
+            };
+            let Some(target_type) = self.resolve_target(to_expr, source_class) else {
+                continue;
+            };
+            let target_class = match &target_type {
+                Type::ClassType(cls_type) => cls_type.class_object(),
+                Type::ClassDef(class_def) => class_def,
+                _ => continue,
+            };
+            if target_class != cls {
+                continue;
+            }
+
+            let related_name =
+                match self.django_related_name(call_expr, source_class, relation_kind) {
+                    Some(name) => name,
+                    None => continue,
+                };
+            let related_type = match self.django_reverse_field_type(relation_kind, source_class) {
+                Some(ty) => ty,
+                None => continue,
+            };
+
+            fields.insert(related_name, ClassSynthesizedField::new(related_type));
+        }
+
+        if fields.is_empty() {
+            None
+        } else {
+            Some(ClassSynthesizedFields::new(fields))
+        }
+    }
+
+    fn django_relation_kind(&self, expr: &Expr) -> Option<DjangoRelationKind> {
+        let ty = self.expr_infer(expr, &self.error_swallower());
+        let field_class = match &ty {
+            Type::ClassType(cls) => cls.class_object(),
+            Type::ClassDef(cls) => cls,
+            _ => return None,
+        };
+
+        if self.is_one_to_one_field(field_class) {
+            Some(DjangoRelationKind::OneToOne)
+        } else if self.is_many_to_many_field(field_class) {
+            Some(DjangoRelationKind::ManyToMany)
+        } else if self.is_foreign_key_field(field_class) {
+            Some(DjangoRelationKind::ForeignKey)
+        } else {
+            None
+        }
+    }
+
+    fn django_reverse_field_type(
+        &self,
+        relation_kind: DjangoRelationKind,
+        source_class: &Class,
+    ) -> Option<Type> {
+        let source_type = self.instantiate(source_class);
+        match relation_kind {
+            DjangoRelationKind::ForeignKey => self.get_related_manager_type(source_type),
+            DjangoRelationKind::ManyToMany => self.get_manager_type(source_type),
+            DjangoRelationKind::OneToOne => Some(source_type),
+        }
+    }
+
+    fn django_related_name(
+        &self,
+        call_expr: &ExprCall,
+        source_class: &Class,
+        relation_kind: DjangoRelationKind,
+    ) -> Option<Name> {
+        let mut related_name_expr = None;
+        for keyword in &call_expr.arguments.keywords {
+            if keyword
+                .arg
+                .as_ref()
+                .is_some_and(|name| name.as_str() == RELATED_NAME.as_str())
+            {
+                related_name_expr = Some(&keyword.value);
+                break;
+            }
+        }
+
+        match related_name_expr {
+            None => Some(self.django_default_related_name(source_class, relation_kind)),
+            Some(Expr::NoneLiteral(_)) => {
+                Some(self.django_default_related_name(source_class, relation_kind))
+            }
+            Some(Expr::StringLiteral(lit)) => {
+                self.format_related_name(lit.value.to_str(), source_class)
+            }
+            _ => None,
+        }
+    }
+
+    fn django_default_related_name(
+        &self,
+        source_class: &Class,
+        relation_kind: DjangoRelationKind,
+    ) -> Name {
+        let mut name = source_class.name().as_str().to_ascii_lowercase();
+        if matches!(
+            relation_kind,
+            DjangoRelationKind::ForeignKey | DjangoRelationKind::ManyToMany
+        ) {
+            name.push_str("_set");
+        }
+        Name::new(name)
+    }
+
+    fn format_related_name(&self, raw: &str, source_class: &Class) -> Option<Name> {
+        let trimmed = raw.trim();
+        if trimmed.is_empty() || trimmed.ends_with('+') {
+            return None;
+        }
+
+        let class_name = source_class.name().as_str().to_ascii_lowercase();
+        let module_name = source_class.module_name();
+        let module_name = module_name.as_str();
+        let mut module_parts = module_name.rsplit('.');
+        let last = module_parts.next().unwrap_or(module_name);
+        let app_label = if last == "models" {
+            module_parts.next().unwrap_or(last)
+        } else {
+            last
+        }
+        .to_ascii_lowercase();
+
+        let substituted = trimmed
+            .replace("%(class)s", &class_name)
+            .replace("%(app_label)s", &app_label);
+
+        if substituted.is_empty() || substituted.ends_with('+') || substituted.contains('%') {
+            return None;
+        }
+
+        Some(Name::new(substituted))
+    }
 }

pyrefly/lib/test/django/foreign_key.rs

@@ -16,6 +16,7 @@ django_testcase!(
 from typing import assert_type
 
 from django.db import models
+from django.db.models.fields.related_descriptors import RelatedManager
 
 class Reporter(models.Model):
     full_name = models.CharField(max_length=70)
@@ -28,6 +29,9 @@ assert_type(article.reporter, Reporter)
 assert_type(article.reporter.full_name, str)
 assert_type(article.reporter_id, int)
 
+reporter = Reporter()
+assert_type(reporter.article_set, RelatedManager[Article])
+
 class B(Article):
     pass
 

pyrefly/lib/test/django/many_to_many.rs

@@ -30,6 +30,9 @@ assert_type(book.authors.all(), QuerySet[Author, Author])
 assert_type(book.authors.filter(name="Bob"), QuerySet[Author, Author])
 assert_type(book.authors.create(name="Alice"), Author) 
 
+author = Author()
+assert_type(author.books, ManyRelatedManager[Book, models.Model])
+
 book.authors.add("wrong type") # E: Argument `Literal['wrong type']` is not assignable to parameter `*objs` with type `Author | int` 
 "#,
 );