fix Support for Django Manager/QuerySet #3648

crates/pyrefly_types/src/class.rs

@@ -110,6 +110,8 @@ pub struct ClassFieldProperties {
     // The field is defined in the class body (not in a method via self.x = ...)
     is_defined_in_class_body: bool,
     attrs_field_specifier: Option<AttrsFieldSpecifier>,
+    /// An explicit Django method-copy policy set in the class body.
+    django_queryset_only: Option<bool>,
     range: TextRange,
     // The range of the docstring following this field, if present
     docstring_range: Option<TextRange>,
@@ -134,6 +136,7 @@ impl ClassFieldProperties {
             is_initialized_on_class: has_default_value,
             is_defined_in_class_body,
             attrs_field_specifier,
+            django_queryset_only: None,
             range,
             docstring_range,
         }
@@ -143,6 +146,11 @@ impl ClassFieldProperties {
         self.is_initialized_on_class
     }
 
+    /// Attach the method-copy policy found during binding.
+    pub fn with_django_queryset_only(mut self, value: Option<bool>) -> Self {
+        self.django_queryset_only = value;
+        self
+    }
     pub fn is_defined_in_class_body(&self) -> bool {
         self.is_defined_in_class_body
     }
@@ -205,6 +213,11 @@ impl ClassFields {
         self.0.get(name).is_some_and(|prop| prop.is_annotated)
     }
 
+    /// Return the explicit Django method-copy policy, if present.
+    pub fn django_queryset_only(&self, name: &Name) -> Option<bool> {
+        self.0.get(name).and_then(|prop| prop.django_queryset_only)
+    }
+
     pub fn is_attrs_field_specifier(&self, name: &Name) -> bool {
         self.0
             .get(name)
@@ -450,7 +463,38 @@ impl Class {
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[derive(Visit, VisitMut, TypeEq)]
-pub struct ClassType(Class, TArgs);
+pub struct ClassType(Class, TArgs, Option<Arc<DjangoClassMetadata>>);
+
+/// Construction-time Django specializations on nominal class instances.
+///
+/// The manager resolves generated members against its queryset. A specialized queryset uses
+/// its effective base for inherited members and nominal upcasts. Type traversal and substitution
+/// include this data, and instance binding preserves it when replacing `Self`. Explicit method
+/// return types stand alone: neither their model nor their row is rebound to the manager's model.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(TypeEq)]
+pub enum DjangoClassMetadata {
+    Manager { queryset: ClassType },
+    QuerySet { base: ClassType },
+}
+
+impl Visit<Type> for Arc<DjangoClassMetadata> {
+    fn recurse<'a>(&'a self, f: &mut dyn FnMut(&'a Type)) {
+        match self.as_ref() {
+            DjangoClassMetadata::Manager { queryset } => queryset.visit(f),
+            DjangoClassMetadata::QuerySet { base } => base.visit(f),
+        }
+    }
+}
+
+impl VisitMut<Type> for Arc<DjangoClassMetadata> {
+    fn recurse_mut(&mut self, f: &mut dyn FnMut(&mut Type)) {
+        match Arc::make_mut(self) {
+            DjangoClassMetadata::Manager { queryset } => queryset.visit_mut(f),
+            DjangoClassMetadata::QuerySet { base } => base.visit_mut(f),
+        }
+    }
+}
 
 impl Display for ClassType {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -462,7 +506,7 @@ impl ClassType {
     /// Create a class type.
     /// The `targs` must match the `tparams`, if this fails we will panic.
     pub fn new(class: Class, targs: TArgs) -> Self {
-        Self(class, targs)
+        Self(class, targs, None)
     }
 
     pub fn class_object(&self) -> &Class {
@@ -490,7 +534,23 @@ impl ClassType {
     /// This is used to propagate instantiation of base class type parameters when computing
     /// the MRO.
     pub fn substitute_with(&self, substitution: &Substitution) -> Self {
-        Self(self.0.dupe(), self.1.substitute_with(substitution))
+        let mut result = self.clone();
+        result.1 = self.1.substitute_with(substitution);
+        if let Some(metadata) = &mut result.2 {
+            metadata.recurse_mut(&mut |ty| substitution.substitute_into_mut(ty));
+        }
+        result
+    }
+
+    /// Attach a Django specialization to this nominal class instance.
+    pub fn with_django_metadata(mut self, metadata: DjangoClassMetadata) -> Self {
+        self.2 = Some(Arc::new(metadata));
+        self
+    }
+
+    /// Return the construction-time Django specialization, if present.
+    pub fn django_metadata(&self) -> Option<&DjangoClassMetadata> {
+        self.2.as_deref()
     }
 
     pub fn substitution(&self) -> Substitution<'_> {

pyrefly/lib/alt/attr.rs

@@ -10,6 +10,7 @@ use std::iter;
 use dupe::Dupe;
 use pyrefly_python::dunder;
 use pyrefly_python::module_name::ModuleName;
+use pyrefly_types::class::DjangoClassMetadata;
 use pyrefly_types::heap::TypeHeap;
 use pyrefly_types::literal::LitEnum;
 use pyrefly_types::shaped_array::ShapedArrayType;
@@ -3504,6 +3505,23 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         res: &mut Vec<AttrInfo>,
     ) {
         self.completions_class(cls.class_object(), expected_attribute_name, res);
+        if let Some(DjangoClassMetadata::Manager { queryset }) = cls.django_metadata() {
+            let mut forwarded = Vec::new();
+            self.completions_class(
+                queryset.class_object(),
+                expected_attribute_name,
+                &mut forwarded,
+            );
+            for info in forwarded {
+                if !res.iter().any(|existing| existing.name == info.name)
+                    && self
+                        .get_django_instance_attribute(cls, &info.name)
+                        .is_some()
+                {
+                    res.push(info);
+                }
+            }
+        }
     }
 
     fn completions_module(

pyrefly/lib/alt/class/class_field.rs

@@ -627,7 +627,7 @@ impl ClassField {
         }
     }
 
-    fn instantiate_helper(&self, f: &mut dyn FnMut(&mut Type)) -> Self {
+    pub(crate) fn instantiate_helper(&self, f: &mut dyn FnMut(&mut Type)) -> Self {
         match &self.0 {
             ClassFieldInner::Property { ty, is_abstract } => {
                 let mut ty = ty.clone();
@@ -2763,6 +2763,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         range: TextRange,
         errors: &ErrorCollector,
     ) -> Option<Type> {
+        let initial_value_expr = match field_definition {
+            ClassFieldDefinition::AssignedInBody { value, .. } => {
+                if let ExprOrBinding::Expr(expr) = value.as_ref() {
+                    Some(expr)
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        };
+
         self.get_enum_class_field_type(
             class,
             name,
@@ -2775,19 +2786,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         )
         .or_else(|| self.get_property_class_field_type(class, name, field_definition))
         .or_else(|| self.get_pydantic_root_model_class_field_type(class, name))
-        .or_else(|| {
-            let initial_value_expr = match field_definition {
-                ClassFieldDefinition::AssignedInBody { value, .. } => {
-                    if let ExprOrBinding::Expr(expr) = value.as_ref() {
-                        Some(expr)
-                    } else {
-                        None
-                    }
-                }
-                _ => None,
-            };
-            self.get_django_field_type(ty, class, Some(name), initial_value_expr)
-        })
+        .or_else(|| self.get_django_field_type(ty, class, Some(name), initial_value_expr))
     }
 
     /// Recognize `x = property(fget, fset, fdel)` and return the corresponding
@@ -3283,6 +3282,25 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         // Otherwise, analyze the value to determine the type
         let (inherited_ty, inherited_annotation) =
             self.get_inherited_type_and_annotation(class, name);
+        if !inferred_from_method
+            && direct_annotation.is_none()
+            && self.get_metadata_for_class(class).is_django_model()
+            && inherited_annotation
+                .as_ref()
+                .and_then(|ann| ann.ty.as_ref())
+                .is_none_or(|ty| {
+                    matches!(ty, Type::ClassType(cls) if
+                        cls.has_qname("django.db.models.manager", "Manager")
+                        || cls.has_qname("django.db.models.manager", "BaseManager"))
+                })
+            && let ExprOrBinding::Expr(e) = value
+            && let Some(manager) = self.get_django_manager_from_queryset_type(class, e)
+        {
+            // Standard manager interfaces allow model rebinding. Explicit custom interfaces
+            // retain their declared type and use ordinary assignment checking below.
+            self.attribute_expr_infer(e, None, name, errors);
+            return (manager, None, IsInherited::Maybe);
+        }
         let mut is_inherited = if inherited_ty.is_none() {
             IsInherited::No
         } else {
@@ -3708,7 +3726,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         None
     }
 
-    fn as_instance_attribute(
+    pub(crate) fn as_instance_attribute(
         &self,
         field_name: &Name,
         field: &ClassField,
@@ -3719,7 +3737,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         // quantified to capture `cls`. Note that `__new__` is the only method that needs this
         // special case because it's not legal to use `Self` in other static methods.
         let self_quantified =
-            if field_name == &dunder::NEW && matches!(instance.kind, InstanceKind::ClassType) {
+            if field_name == &dunder::NEW && matches!(instance.kind, InstanceKind::ClassType(_)) {
                 Some(self.get_self_quantified(instance.class, instance.to_type(self.heap)))
             } else {
                 None
@@ -4428,6 +4446,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     continue;
                 }
             }
+            if metadata.is_django_model()
+                && !class_field.has_explicit_annotation()
+                && self.is_django_manager_type(&class_field.ty())
+                && self.is_django_manager_type(&want_class_field.ty())
+            {
+                // Django binds each manager replacement to the receiving model.
+                continue;
+            }
             let want_is_static = want_class_field.is_staticmethod();
             let want_is_class = want_class_field.is_classmethod();
             if (want_is_static || want_is_class) && class_field.is_ordinary_instance_method() {
@@ -5115,6 +5141,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     }
 
     pub fn get_instance_attribute(&self, cls: &ClassType, name: &Name) -> Option<ClassAttribute> {
+        if let Some(attribute) = self.get_django_instance_attribute(cls, name) {
+            return Some(attribute);
+        }
         self.get_class_member(cls.class_object(), name)
             .map(|field| self.as_instance_attribute(name, field.as_ref(), &Instance::of_class(cls)))
     }

pyrefly/lib/alt/class/classdef.rs

@@ -10,6 +10,7 @@ use std::sync::Arc;
 use dupe::Dupe;
 use pyrefly_python::nesting_context::NestingContext;
 use pyrefly_types::callable::Callable;
+use pyrefly_types::class::DjangoClassMetadata;
 use pyrefly_types::class::PrecomputedTParams;
 use pyrefly_types::quantified::Quantified;
 use pyrefly_types::special_form::SpecialForm;
@@ -258,6 +259,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     pub fn as_superclass(&self, class: &ClassType, want: &Class) -> Option<ClassType> {
         if class.class_object() == want {
             Some(class.clone())
+        } else if let Some(DjangoClassMetadata::QuerySet { base }) = class.django_metadata()
+            && self.has_superclass(base.class_object(), want)
+        {
+            self.as_superclass(base, want)
         } else if !class.class_object().is_builtin("tuple")
             && let Some(tuple) = self.as_tuple(class)
             && self.has_superclass(self.stdlib.tuple_object(), want)