From b442bf63470d01943232beb00398487628cb59ef Mon Sep 17 00:00:00 2001 From: Gutti Torvadal Date: Sun, 18 Oct 2026 23:52:34 +0200 Subject: [PATCH 1/2] Add failing test for schema annotations on Polars lazy frames --- pyrefly/lib/test/polars/dataframe.rs | 2 ++ pyrefly/lib/test/polars/third-party/polars/__init__.pyi | 1 + 2 files changed, 3 insertions(+) diff --git a/pyrefly/lib/test/polars/dataframe.rs b/pyrefly/lib/test/polars/dataframe.rs index d1fdcf80fb..03e94ddac3 100644 --- a/pyrefly/lib/test/polars/dataframe.rs +++ b/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] "#, ); diff --git a/pyrefly/lib/test/polars/third-party/polars/__init__.pyi b/pyrefly/lib/test/polars/third-party/polars/__init__.pyi index f5d71e28ec..8c824021ef 100644 --- a/pyrefly/lib/test/polars/third-party/polars/__init__.pyi +++ b/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 From 29069c90dd0bc17a256d7f108147a88730b7ab83 Mon Sep 17 00:00:00 2001 From: Gutti Torvadal Date: Sun, 18 Oct 2026 23:55:06 +0200 Subject: [PATCH 2/2] Typecheck annotated Polars lazy frames --- pyrefly/lib/alt/polars_specials.rs | 8 +++++--- pyrefly/lib/alt/solve.rs | 5 ++--- pyrefly/lib/alt/specials.rs | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/pyrefly/lib/alt/polars_specials.rs b/pyrefly/lib/alt/polars_specials.rs index eb6cfa1cb4..6fc14e6e1d 100644 --- a/pyrefly/lib/alt/polars_specials.rs +++ b/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 { + /// 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 { 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 { diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 43c9bf8a10..9dd12e95fd 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/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() { diff --git a/pyrefly/lib/alt/specials.rs b/pyrefly/lib/alt/specials.rs index f1d0bb9a56..6917d69c4c 100644 --- a/pyrefly/lib/alt/specials.rs +++ b/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()) }