Support annotated Polars lazy frames

pyrefly/lib/alt/polars_specials.rs

@@ -1656,12 +1656,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         Some(entries)
     }
 
-    /// Interpret supported `Annotated[pl.DataFrame, ...]` metadata as a schema contract.
-    pub fn polars_dataframe_annotated_type(&self, inner: &Type, metadata: &[Type]) -> Option<Type> {
+    /// Interpret supported `Annotated[pl.DataFrame, ...]` or `Annotated[pl.LazyFrame, ...]`
+    /// metadata as a schema contract.
+    pub fn polars_frame_annotated_type(&self, inner: &Type, metadata: &[Type]) -> Option<Type> {
         let Type::ClassType(underlying) = inner else {
             return None;
         };
-        if !is_polars_dataframe(underlying.class_object()) {
+        let class = underlying.class_object();
+        if !is_polars_dataframe(class) && !is_polars_lazyframe(class) {
             return None;
         }
         let (schema_cls, completeness) = match metadata {

pyrefly/lib/alt/solve.rs

@@ -1023,10 +1023,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                         .map(|e| self.expr_infer(e, &self.error_swallower()))
                         .collect();
                     if let Some(inner) = ann.ty.as_ref()
-                        && let Some(dataframe) =
-                            self.polars_dataframe_annotated_type(inner, &metadata)
+                        && let Some(frame) = self.polars_frame_annotated_type(inner, &metadata)
                     {
-                        ann.ty = Some(dataframe);
+                        ann.ty = Some(frame);
                     }
                 }
                 if qualifier == Qualifier::ClassVar && ann.get_type().contains_type_variable() {

pyrefly/lib/alt/specials.rs

@@ -595,8 +595,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     .iter()
                     .map(|e| self.expr_infer(e, &self.error_swallower()))
                     .collect();
-                if let Some(dataframe) = self.polars_dataframe_annotated_type(&inner, &metadata) {
-                    inner = dataframe;
+                if let Some(frame) = self.polars_frame_annotated_type(&inner, &metadata) {
+                    inner = frame;
                 }
                 Type::Annotated(Box::new(inner), metadata.into_boxed_slice())
             }

pyrefly/lib/test/polars/dataframe.rs

@@ -5564,6 +5564,8 @@ class MySchema:
     asset: pl.String
 def f(df: Annotated[pl.DataFrame, MySchema]) -> None:
     reveal_type(df)  # E: revealed type: DataFrame[price: Float64, asset: String]
+def g(lf: Annotated[pl.LazyFrame, MySchema]) -> None:
+    reveal_type(lf)  # E: revealed type: LazyFrame[price: Float64, asset: String]
 "#,
 );
 

pyrefly/lib/test/polars/third-party/polars/__init__.pyi

@@ -8,6 +8,7 @@ from polars.functions.len import len as len
 from polars.functions.lit import lit as lit
 from polars.functions.whenthen import when as when
 from polars.io.csv.functions import read_csv as read_csv, scan_csv as scan_csv
+from polars.lazyframe.frame import LazyFrame as LazyFrame
 from polars.schema import Schema as Schema
 from polars.series.series import Series as Series