generate TypeFormContext from AnnotationTarget

yangdanny97 · facebook-github-bot · commit 97b80e38ce1e · 2025-03-11T21:19:42.000-07:00
Summary: As suggested by Rebecca, this avoids some duplication in the annotation binding/solving logic. The more granular AnnotationTarget will also allow us to better model *args and *kwargs types, to help us fix facebook/pyrefly#16 (comment) Reviewed By: rchen152 Differential Revision: D71012503 fbshipit-source-id: 4015d42dcd520821f9de21421f4f8c3b0d68d10e
diff --git a/pyre2/pyre2/lib/alt/solve.rs b/pyre2/pyre2/lib/alt/solve.rs
@@ -195,8 +195,9 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         errors: &ErrorCollector,
     ) -> Arc<AnnotationWithTarget> {
         match binding {
-            BindingAnnotation::AnnotateExpr(target, x, self_type, ctx) => {
-                let mut ann = self.expr_annotation(x, *ctx, errors);
+            BindingAnnotation::AnnotateExpr(target, x, self_type) => {
+                let type_form_context = target.type_form_context();
+                let mut ann = self.expr_annotation(x, type_form_context, errors);
                 if let Some(self_type) = self_type
                     && let Some(ty) = &mut ann.ty
                 {
diff --git a/pyre2/pyre2/lib/binding/binding.rs b/pyre2/pyre2/lib/binding/binding.rs
@@ -942,6 +942,8 @@ impl Display for AnnotationWithTarget {
 pub enum AnnotationTarget {
     /// A function parameter with a type annotation
     Param(Name),
+    ArgsParam(Name),
+    KwargsParam(Name),
     /// A return type annotation on a function. The name is that of the function
     Return(Name),
     /// An annotated assignment. For attribute assignments, the name is the attribute name ("attr" in "x.attr")
@@ -954,32 +956,42 @@ impl Display for AnnotationTarget {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::Param(name) => write!(f, "param {name}"),
+            Self::ArgsParam(name) => write!(f, "args {name}"),
+            Self::KwargsParam(name) => write!(f, "kwargs {name}"),
             Self::Return(name) => write!(f, "{name} return"),
             Self::Assign(name) => write!(f, "var {name}"),
             Self::ClassMember(name) => write!(f, "attr {name}"),
         }
     }
 }
 
+impl AnnotationTarget {
+    pub fn type_form_context(&self) -> TypeFormContext {
+        match self {
+            Self::Param(_) => TypeFormContext::ParameterAnnotation,
+            Self::ArgsParam(_) => TypeFormContext::ParameterArgsAnnotation,
+            Self::KwargsParam(_) => TypeFormContext::ParameterKwargsAnnotation,
+            Self::Return(_) => TypeFormContext::ReturnAnnotation,
+            Self::Assign(_) => TypeFormContext::VarAnnotation,
+            Self::ClassMember(_) => TypeFormContext::ClassVarAnnotation,
+        }
+    }
+}
+
 /// Values that return an annotation.
 #[derive(Clone, Debug)]
 pub enum BindingAnnotation {
     /// The type is annotated to be this key, will have the outer type removed.
     /// Optionally occuring within a class, in which case Self refers to this class.
-    AnnotateExpr(
-        AnnotationTarget,
-        Expr,
-        Option<Idx<KeyClass>>,
-        TypeFormContext,
-    ),
+    AnnotateExpr(AnnotationTarget, Expr, Option<Idx<KeyClass>>),
     /// A literal type we know statically.
     Type(AnnotationTarget, Type),
 }
 
 impl DisplayWith<Bindings> for BindingAnnotation {
     fn fmt(&self, f: &mut fmt::Formatter<'_>, ctx: &Bindings) -> fmt::Result {
         match self {
-            Self::AnnotateExpr(_, x, self_type, _) => write!(
+            Self::AnnotateExpr(_, x, self_type) => write!(
                 f,
                 "_: {}{}",
                 ctx.module_info().display(x),
diff --git a/pyre2/pyre2/lib/binding/bindings.rs b/pyre2/pyre2/lib/binding/bindings.rs
@@ -33,7 +33,6 @@ use starlark_map::small_map::SmallMap;
 use starlark_map::small_set::SmallSet;
 use vec1::Vec1;
 
-use crate::alt::solve::TypeFormContext;
 use crate::binding::binding::AnnotationTarget;
 use crate::binding::binding::Binding;
 use crate::binding::binding::BindingAnnotation;
@@ -673,21 +672,16 @@ impl<'a> BindingsBuilder<'a> {
 
     pub fn bind_function_param(
         &mut self,
+        target: AnnotationTarget,
         x: AnyParameterRef,
         function_idx: Idx<KeyFunction>,
         self_type: Option<Idx<KeyClass>>,
-        ctx: TypeFormContext,
     ) {
         let name = x.name();
         let annot = x.annotation().map(|x| {
             self.table.insert(
                 KeyAnnotation::Annotation(ShortIdentifier::new(name)),
-                BindingAnnotation::AnnotateExpr(
-                    AnnotationTarget::Param(name.id.clone()),
-                    x.clone(),
-                    self_type,
-                    ctx,
-                ),
+                BindingAnnotation::AnnotateExpr(target.clone(), x.clone(), self_type),
             )
         });
         let (annot, def) = match annot {
@@ -696,10 +690,7 @@ impl<'a> BindingsBuilder<'a> {
                 let var = self.solver.fresh_contained(self.uniques);
                 let annot = self.table.insert(
                     KeyAnnotation::Annotation(ShortIdentifier::new(name)),
-                    BindingAnnotation::Type(
-                        AnnotationTarget::Param(name.id.clone()),
-                        var.to_type(),
-                    ),
+                    BindingAnnotation::Type(target, var.to_type()),
                 );
                 (annot, Either::Right((var, function_idx)))
             }
diff --git a/pyre2/pyre2/lib/binding/class.rs b/pyre2/pyre2/lib/binding/class.rs
@@ -22,7 +22,6 @@ use ruff_text_size::TextRange;
 use starlark_map::small_map::SmallMap;
 
 use crate::alt::class::class_metadata::BaseClass;
-use crate::alt::solve::TypeFormContext;
 use crate::binding::binding::AnnotationTarget;
 use crate::binding::binding::Binding;
 use crate::binding::binding::BindingAnnotation;
@@ -360,7 +359,6 @@ impl<'a> BindingsBuilder<'a> {
                         AnnotationTarget::ClassMember(member_name.clone()),
                         annotation,
                         None,
-                        TypeFormContext::ClassVarAnnotation,
                     )
                 };
                 Some(self.table.insert(ann_key, ann_val))
diff --git a/pyre2/pyre2/lib/binding/function.rs b/pyre2/pyre2/lib/binding/function.rs
@@ -17,7 +17,6 @@ use ruff_python_ast::StmtExpr;
 use ruff_python_ast::StmtFunctionDef;
 use ruff_text_size::Ranged;
 
-use crate::alt::solve::TypeFormContext;
 use crate::ast::Ast;
 use crate::binding::binding::AnnotationTarget;
 use crate::binding::binding::Binding;
@@ -61,26 +60,26 @@ impl<'a> BindingsBuilder<'a> {
                 self_name = Some(x.parameter.name.clone());
             }
             self.bind_function_param(
+                AnnotationTarget::Param(x.parameter.name.id.clone()),
                 AnyParameterRef::NonVariadic(x),
                 function_idx,
                 self_type,
-                TypeFormContext::ParameterAnnotation,
             );
         }
         if let Some(box args) = &x.vararg {
             self.bind_function_param(
+                AnnotationTarget::ArgsParam(args.name.id.clone()),
                 AnyParameterRef::Variadic(args),
                 function_idx,
                 self_type,
-                TypeFormContext::ParameterArgsAnnotation,
             );
         }
         if let Some(box kwargs) = &x.kwarg {
             self.bind_function_param(
+                AnnotationTarget::KwargsParam(kwargs.name.id.clone()),
                 AnyParameterRef::Variadic(kwargs),
                 function_idx,
                 self_type,
-                TypeFormContext::ParameterKwargsAnnotation,
             );
         }
         if let Scope {
@@ -155,7 +154,6 @@ impl<'a> BindingsBuilder<'a> {
                         AnnotationTarget::Return(func_name.id.clone()),
                         *x,
                         self_type,
-                        TypeFormContext::ReturnAnnotation,
                     ),
                 ),
             )
diff --git a/pyre2/pyre2/lib/binding/stmt.rs b/pyre2/pyre2/lib/binding/stmt.rs
@@ -18,7 +18,6 @@ use ruff_python_ast::StmtImportFrom;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 
-use crate::alt::solve::TypeFormContext;
 use crate::ast::Ast;
 use crate::binding::binding::AnnotationStyle;
 use crate::binding::binding::AnnotationTarget;
@@ -345,14 +344,13 @@ impl<'a> BindingsBuilder<'a> {
                         )
                     } else {
                         BindingAnnotation::AnnotateExpr(
-                            AnnotationTarget::Assign(name.id.clone()),
-                            *x.annotation.clone(),
-                            None,
                             if in_class_body {
-                                TypeFormContext::ClassVarAnnotation
+                                AnnotationTarget::ClassMember(name.id.clone())
                             } else {
-                                TypeFormContext::VarAnnotation
+                                AnnotationTarget::Assign(name.id.clone())
                             },
+                            *x.annotation.clone(),
+                            None,
                         )
                     };
                     let ann_key = self.table.insert(ann_key, ann_val);
@@ -410,10 +408,9 @@ impl<'a> BindingsBuilder<'a> {
                     let ann_key = self.table.insert(
                         KeyAnnotation::AttrAnnotation(x.annotation.range()),
                         BindingAnnotation::AnnotateExpr(
-                            AnnotationTarget::Assign(attr.attr.id.clone()),
+                            AnnotationTarget::ClassMember(attr.attr.id.clone()),
                             *x.annotation,
                             None,
-                            TypeFormContext::ClassVarAnnotation,
                         ),
                     );
                     let value_binding = match &x.value {
diff --git a/pyre2/pyre2/lib/error/context.rs b/pyre2/pyre2/lib/error/context.rs
@@ -112,7 +112,9 @@ pub enum TypeCheckKind {
 impl TypeCheckKind {
     pub fn from_annotation_target(target: &AnnotationTarget) -> Self {
         match target {
-            AnnotationTarget::Param(_) => Self::Unknown,
+            AnnotationTarget::Param(_)
+            | AnnotationTarget::ArgsParam(_)
+            | AnnotationTarget::KwargsParam(_) => Self::Unknown,
             AnnotationTarget::Return(_func) => Self::ExplicitFunctionReturn,
             AnnotationTarget::Assign(name) => Self::AnnotatedName(name.clone()),
             AnnotationTarget::ClassMember(member) => Self::Attribute(member.clone()),
