Correctly synthesize pk and id types based on custom pk info

Reviewed By: rchen152

Differential Revision: D85589929

fbshipit-source-id: 820238d8c8eabe7e27fdf6430d44f5251ab847ba

Author: migeed-z

pyrefly/lib/alt/class/django.rs

@@ -197,14 +197,31 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
 
         let mut fields = SmallMap::new();
 
-        let auto_field_export = KeyExport(AUTO_FIELD);
-        let auto_field_type =
-            self.get_from_export(ModuleName::django_models_fields(), None, &auto_field_export);
-
-        // TODO: Extend the solution to handle custom pk
-        if let Some(id_type) = self.get_django_field_type(&auto_field_type, cls) {
-            fields.insert(ID, ClassSynthesizedField::new(id_type.clone()));
-            fields.insert(PK, ClassSynthesizedField::new(id_type));
+        let custom_pk_field = metadata
+            .django_model_metadata()
+            .and_then(|dm| dm.custom_primary_key_field.as_ref());
+
+        if let Some(pk_field_name) = custom_pk_field {
+            let instance_type = self.as_class_type_unchecked(cls).to_type();
+            let pk_attr_type = self.attr_infer_for_type(
+                &instance_type,
+                pk_field_name,
+                TextRange::default(),
+                &self.error_swallower(),
+                None,
+            );
+            fields.insert(PK, ClassSynthesizedField::new(pk_attr_type));
+            // When there's a custom pk, don't synthesize an `id` field
+        } else {
+            // No custom pk, use default AutoField for both id and pk
+            let auto_field_export = KeyExport(AUTO_FIELD);
+            let auto_field_type =
+                self.get_from_export(ModuleName::django_models_fields(), None, &auto_field_export);
+
+            if let Some(id_type) = self.get_django_field_type(&auto_field_type, cls) {
+                fields.insert(ID, ClassSynthesizedField::new(id_type.clone()));
+                fields.insert(PK, ClassSynthesizedField::new(id_type));
+            }
         }
 
         Some(ClassSynthesizedFields::new(fields))

pyrefly/lib/test/django/auto_generated_fields.rs

@@ -39,7 +39,6 @@ assert_type(reporter.id, str)
 );
 
 django_testcase!(
-    bug = "do not generate an id field when a custom key is present; wrong pk type",
     test_custom_pk,
     r#"
 from typing import assert_type
@@ -53,13 +52,13 @@ class B(Article):
     pass
 
 article = Article()
-article.id
+article.id # E: Object of class `Article` has no attribute `id`
 assert_type(article.uuid, UUID) 
-assert_type(article.pk, UUID) # E: assert_type(int, UUID) 
+assert_type(article.pk, UUID) 
 
 article2 = B()
-article2.id
+article2.id # E: Object of class `B` has no attribute `id` 
 assert_type(article2.uuid, UUID)
-assert_type(article2.pk, UUID) # E: assert_type(int, UUID) 
+assert_type(article2.pk, UUID) 
 "#,
 );