From 217994abd57c850cc0f47bdd07e952b63d4a7bd9 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:26:26 +0100 Subject: [PATCH] Narwhals support --- .cargo/config.toml | 1 + crates/pyrefly_types/src/data_frame.rs | 7 + pyrefly/lib/alt/expr.rs | 9 +- pyrefly/lib/alt/polars_specials.rs | 312 ++++++++++++++---- pyrefly/lib/solver/subset.rs | 7 +- pyrefly/lib/state/lsp/dict_completions.rs | 33 +- pyrefly/lib/test/mod.rs | 1 + pyrefly/lib/test/narwhals.rs | 10 + pyrefly/lib/test/narwhals/dataframe.rs | 306 +++++++++++++++++ .../third-party/narwhals/__init__.pyi | 30 ++ .../third-party/narwhals/dataframe.pyi | 75 +++++ .../narwhals/third-party/narwhals/dtypes.pyi | 23 ++ .../narwhals/third-party/narwhals/expr.pyi | 34 ++ .../third-party/narwhals/functions.pyi | 17 + .../narwhals/third-party/narwhals/schema.pyi | 2 + .../narwhals/third-party/narwhals/series.pyi | 9 + .../third-party/narwhals/stable/__init__.pyi | 0 .../narwhals/stable/v1/__init__.pyi | 38 +++ .../narwhals/stable/v2/__init__.pyi | 38 +++ pyrefly/lib/test/narwhals/util.rs | 49 +++ website/docs/dataframes.mdx | 86 ++++- 21 files changed, 970 insertions(+), 117 deletions(-) create mode 100644 pyrefly/lib/test/narwhals.rs create mode 100644 pyrefly/lib/test/narwhals/dataframe.rs create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/__init__.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/dataframe.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/dtypes.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/expr.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/functions.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/schema.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/series.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/stable/__init__.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/stable/v1/__init__.pyi create mode 100644 pyrefly/lib/test/narwhals/third-party/narwhals/stable/v2/__init__.pyi create mode 100644 pyrefly/lib/test/narwhals/util.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index 140eef8fd5..d50b3e9dac 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -10,6 +10,7 @@ ATTRS_TEST_PATH = { value = "pyrefly/lib/test/attrs/third-party", relative = tru CINDERX_FIXTURES_PATH = { value = "pyrefly/lib/test/cinderx/fixtures", relative = true } MARSHMALLOW_TEST_PATH = { value = "pyrefly/lib/test/marshmallow/third-party", relative = true } POLARS_TEST_PATH = { value = "pyrefly/lib/test/polars/third-party", relative = true } +NARWHALS_TEST_PATH = { value = "pyrefly/lib/test/narwhals/third-party", relative = true } GLEAN_SNAPSHOTS_PATH = { value = "pyrefly/lib/report/glean/snapshots", relative = true } COVERAGE_TEST_PATH = { value = "pyrefly/lib/test/coverage/test_files", relative = true } STUBGEN_TEST_PATH = { value = "pyrefly/lib/test/stubgen", relative = true } diff --git a/crates/pyrefly_types/src/data_frame.rs b/crates/pyrefly_types/src/data_frame.rs index ad3311a146..f406334c8b 100644 --- a/crates/pyrefly_types/src/data_frame.rs +++ b/crates/pyrefly_types/src/data_frame.rs @@ -39,9 +39,16 @@ impl SchemaCompleteness { )] pub enum DataFrameKind { Polars, + Narwhals, Pandas, } +impl DataFrameKind { + pub fn is_polars_api(self) -> bool { + matches!(self, Self::Polars | Self::Narwhals) + } +} + /// Whether a schema is inferred or declared by an explicit annotation. #[derive( Debug, PartialOrd, Ord, Clone, Copy, Eq, PartialEq, Hash, Visit, VisitMut, TypeEq diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index df10b72eec..5d2a394a8a 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/pyrefly/lib/alt/expr.rs @@ -19,7 +19,6 @@ use pyrefly_python::dunder; use pyrefly_python::module_name::ModuleName; use pyrefly_python::nesting_context::NestingContext; use pyrefly_python::short_identifier::ShortIdentifier; -use pyrefly_types::data_frame::DataFrameKind; use pyrefly_types::dimension::Int; use pyrefly_types::dimension::canonicalize; use pyrefly_types::dimension::gradual_size; @@ -94,7 +93,7 @@ use crate::alt::call::CallStyle; use crate::alt::callable::CallArg; use crate::alt::class::typed_dict::TypedDictErrorKind; use crate::alt::nn_module_specials::is_nn_module_dict; -use crate::alt::polars_specials::is_polars_series; +use crate::alt::polars_specials::is_polars_api_series; use crate::alt::regex::RegexValidationError; use crate::alt::regex::validate_pattern; use crate::alt::shape_extension::is_int_tuple_bound; @@ -4441,7 +4440,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } Type::DataFrame(schema) => { if let Expr::List(ExprList { elts, .. }) = slice - && schema.kind == DataFrameKind::Polars + && schema.kind.is_polars_api() { if elts.is_empty() { return Type::DataFrame(schema); @@ -4457,7 +4456,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { && let Some(name) = self.polars_column_name(slice) { match schema.columns.iter().find(|(c, _)| **c == name) { - Some((_, dtype)) if schema.kind == DataFrameKind::Polars => { + Some((_, dtype)) if schema.kind.is_polars_api() => { column_dtype = Some(dtype.clone()); } Some(_) => {} @@ -4484,7 +4483,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // Preserve the stub's Series class when attaching an element dtype. match (column_dtype, result) { (Some(dtype), Type::ClassType(cls)) - if is_polars_series(cls.class_object()) => + if is_polars_api_series(cls.class_object()) => { SeriesSchema { underlying: cls, diff --git a/pyrefly/lib/alt/polars_specials.rs b/pyrefly/lib/alt/polars_specials.rs index 6fc14e6e1d..4824ee7d75 100644 --- a/pyrefly/lib/alt/polars_specials.rs +++ b/pyrefly/lib/alt/polars_specials.rs @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -//! Column-aware typing for Polars and pandas DataFrames. +//! Column-aware typing for Polars, Narwhals, and pandas DataFrames. use pyrefly_types::data_frame::DataFrameKind; use pyrefly_types::data_frame::DataFrameSchema; @@ -52,6 +52,9 @@ use crate::types::literal::Lit; const POLARS_MODULE: &str = "polars"; const POLARS_MODULE_PREFIX: &str = "polars."; +const NARWHALS_MODULE: &str = "narwhals"; +const NARWHALS_MODULE_PREFIX: &str = "narwhals."; +const NARWHALS_STABLE_NAMESPACES: [&str; 2] = ["narwhals.stable.v1", "narwhals.stable.v2"]; const POLARS_ALL_COLUMNS: &str = "*"; const POLARS_DEFAULT_INFER_SCHEMA_LENGTH: usize = 100; const POLARS_DEFAULT_JOIN_SUFFIX: &str = "_right"; @@ -67,6 +70,10 @@ enum RuntimeClass { PolarsLazyFrame, PolarsCol, PolarsSchema, + NarwhalsDataFrame, + NarwhalsLazyFrame, + NarwhalsSeries, + NarwhalsExpr, PandasDataFrame, Date, Datetime, @@ -75,22 +82,47 @@ enum RuntimeClass { } impl RuntimeClass { + fn qnames(self) -> &'static [(&'static str, &'static str)] { + match self { + Self::PolarsDataFrame => &[("polars.dataframe.frame", "DataFrame")], + Self::PolarsSeries => &[("polars.series.series", "Series")], + Self::PolarsDataFrameSeries => &[("polars.dataframe.frame", "Series")], + Self::PolarsExpr => &[("polars.expr.expr", "Expr")], + Self::PolarsLazyFrame => &[("polars.lazyframe.frame", "LazyFrame")], + Self::PolarsCol => &[("polars.functions.col", "Col")], + Self::PolarsSchema => &[("polars.schema", "Schema")], + Self::NarwhalsDataFrame => &[ + ("narwhals.dataframe", "DataFrame"), + ("narwhals.stable.v1", "DataFrame"), + ("narwhals.stable.v2", "DataFrame"), + ], + Self::NarwhalsLazyFrame => &[ + ("narwhals.dataframe", "LazyFrame"), + ("narwhals.stable.v1", "LazyFrame"), + ("narwhals.stable.v2", "LazyFrame"), + ], + Self::NarwhalsSeries => &[ + ("narwhals.series", "Series"), + ("narwhals.stable.v1", "Series"), + ("narwhals.stable.v2", "Series"), + ], + Self::NarwhalsExpr => &[ + ("narwhals.expr", "Expr"), + ("narwhals.stable.v1", "Expr"), + ("narwhals.stable.v2", "Expr"), + ], + Self::PandasDataFrame => &[("pandas.core.frame", "DataFrame")], + Self::Date => &[("datetime", "date")], + Self::Datetime => &[("datetime", "datetime")], + Self::Time => &[("datetime", "time")], + Self::Timedelta => &[("datetime", "timedelta")], + } + } + fn matches(self, cls: &Class) -> bool { - let (module, name) = match self { - Self::PolarsDataFrame => ("polars.dataframe.frame", "DataFrame"), - Self::PolarsSeries => ("polars.series.series", "Series"), - Self::PolarsDataFrameSeries => ("polars.dataframe.frame", "Series"), - Self::PolarsExpr => ("polars.expr.expr", "Expr"), - Self::PolarsLazyFrame => ("polars.lazyframe.frame", "LazyFrame"), - Self::PolarsCol => ("polars.functions.col", "Col"), - Self::PolarsSchema => ("polars.schema", "Schema"), - Self::PandasDataFrame => ("pandas.core.frame", "DataFrame"), - Self::Date => ("datetime", "date"), - Self::Datetime => ("datetime", "datetime"), - Self::Time => ("datetime", "time"), - Self::Timedelta => ("datetime", "timedelta"), - }; - cls.has_toplevel_qname(module, name) + self.qnames() + .iter() + .any(|(module, name)| cls.has_toplevel_qname(module, name)) } } @@ -98,8 +130,8 @@ fn is_polars_dataframe(cls: &Class) -> bool { RuntimeClass::PolarsDataFrame.matches(cls) } -pub fn is_polars_series(cls: &Class) -> bool { - RuntimeClass::PolarsSeries.matches(cls) +pub fn is_polars_api_series(cls: &Class) -> bool { + RuntimeClass::PolarsSeries.matches(cls) || RuntimeClass::NarwhalsSeries.matches(cls) } /// Identifies the callable object that also supports `pl.col.name` access. @@ -107,14 +139,30 @@ pub fn is_polars_col(cls: &Class) -> bool { RuntimeClass::PolarsCol.matches(cls) } -fn is_polars_expr(cls: &Class) -> bool { - RuntimeClass::PolarsExpr.matches(cls) +fn is_polars_api_expr(cls: &Class) -> bool { + RuntimeClass::PolarsExpr.matches(cls) || RuntimeClass::NarwhalsExpr.matches(cls) } fn is_polars_lazyframe(cls: &Class) -> bool { RuntimeClass::PolarsLazyFrame.matches(cls) } +/// The library an eager frame class belongs to, for the libraries following the Polars API. +fn eager_frame_kind(cls: &Class) -> Option { + if is_polars_dataframe(cls) { + Some(DataFrameKind::Polars) + } else if RuntimeClass::NarwhalsDataFrame.matches(cls) { + Some(DataFrameKind::Narwhals) + } else { + None + } +} + +/// Whether the class is a lazy frame of a library following the Polars API. +fn is_lazy_frame(cls: &Class) -> bool { + is_polars_lazyframe(cls) || RuntimeClass::NarwhalsLazyFrame.matches(cls) +} + fn column_transform_schema<'b>(base: &'b Type, args: &Arguments) -> Option<&'b DataFrameSchema> { let Type::DataFrame(schema) = base else { return None; @@ -236,7 +284,8 @@ fn is_polars_selector_name(name: &Name) -> bool { } /// Map a resolved type to the Polars dtype it names, e.g. the `pl.Float64` class to `Float64`. -/// Only the modeled scalar dtypes from the `polars` package are recognized; anything else is `None`. +/// Only the modeled scalar dtypes are recognized, and only from a package whose dtype classes +/// carry Polars' names and meanings; anything else is `None`. fn polars_dtype_from_type(ty: &Type) -> Option { let cls = match ty { Type::ClassDef(cls) => cls, @@ -244,7 +293,11 @@ fn polars_dtype_from_type(ty: &Type) -> Option { _ => return None, }; let module = cls.module_name(); - if module.as_str() != POLARS_MODULE && !module.as_str().starts_with(POLARS_MODULE_PREFIX) { + let module = module.as_str(); + let in_package = |root: &str, prefix: &str| module == root || module.starts_with(prefix); + if !in_package(POLARS_MODULE, POLARS_MODULE_PREFIX) + && !in_package(NARWHALS_MODULE, NARWHALS_MODULE_PREFIX) + { return None; } PolarsDType::from_polars_name(cls.name().as_str()) @@ -329,7 +382,30 @@ enum PolarsFrameConversion { impl PolarsFunction { fn from_id(id: &FuncDefId) -> Self { - match (id.qname.id().as_str(), id.qname.module_name().as_str()) { + let name = id.qname.id().as_str(); + let module = id.qname.module_name(); + let module = module.as_str(); + if module == "narwhals.functions" { + return match name { + "col" => Self::Col, + "concat" => Self::Concat, + "len_" => Self::Len, + "lit" => Self::Lit, + "when" => Self::When, + _ => Self::Unmodeled, + }; + } + if NARWHALS_STABLE_NAMESPACES.contains(&module) { + return match name { + "col" => Self::Col, + "concat" => Self::Concat, + "len" => Self::Len, + "lit" => Self::Lit, + "when" => Self::When, + _ => Self::Unmodeled, + }; + } + match (name, module) { ("col", "polars.functions.col") => Self::Col, ("concat", "polars.functions.eager") => Self::Concat, ("len", "polars.functions.len") => Self::Len, @@ -723,6 +799,54 @@ impl ExprValue { ExprValue::Dtype(d) => d.is_integer(), } } + + /// The dtype of the column this value produces. On Narwhals frames, an integer literal that + /// fits in 32 bits is `Unknown` because Narwhals does not guarantee its dtype. + fn column_dtype(self, kind: DataFrameKind) -> PolarsDType { + match self { + ExprValue::IntLit(v) if kind == DataFrameKind::Narwhals && i32::try_from(v).is_ok() => { + PolarsDType::Unknown + } + value => value.dtype(), + } + } +} + +/// Computes the result of an expression node from its operands. On Narwhals frames, the result +/// is `Unknown` unless Narwhals guarantees its dtype, which requires every numeric operand and the +/// result to be `Int64` or `Float64`. `narwhals_unspecified` marks an operation whose dtype +/// Narwhals does not guarantee even then. +fn computed_value( + kind: DataFrameKind, + operands: [ExprValue; N], + narwhals_unspecified: bool, + compute: impl FnOnce([ExprValue; N]) -> Option, +) -> Option { + match kind { + DataFrameKind::Polars => return compute(operands), + DataFrameKind::Narwhals => {} + DataFrameKind::Pandas => { + unreachable!("Polars expressions are only evaluated on Polars-API frames") + } + } + if narwhals_unspecified { + return Some(ExprValue::Dtype(PolarsDType::Unknown)); + } + let unspecified = |value: &ExprValue| match value { + ExprValue::Dtype(PolarsDType::Int64 | PolarsDType::Float64) => false, + ExprValue::Dtype(PolarsDType::Unknown) => true, + ExprValue::Dtype(d) => d.is_numeric(), + ExprValue::IntLit(_) | ExprValue::FloatLit => false, + }; + if operands.iter().any(unspecified) { + return Some(ExprValue::Dtype(PolarsDType::Unknown)); + } + let result = compute(operands)?; + Some(if unspecified(&result) { + ExprValue::Dtype(PolarsDType::Unknown) + } else { + result + }) } fn literal_value(expr: &Expr) -> Option { @@ -747,7 +871,7 @@ fn series_method_schema(base: &Type) -> Option<&DataFrameSchema> { let Type::DataFrame(schema) = base else { return None; }; - (schema.kind == DataFrameKind::Polars && schema.is_complete()).then_some(&**schema) + (schema.kind.is_polars_api() && schema.is_complete()).then_some(&**schema) } fn get_column_name_arg(args: &Arguments) -> Option<&Expr> { @@ -1108,9 +1232,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return self .infer_polars_concat(arguments) .map( - |(columns, completeness)| PolarsCallSpecialization::DataFrame { + |(columns, completeness, kind)| PolarsCallSpecialization::DataFrame { columns, - kind: DataFrameKind::Polars, + kind, completeness, }, ); @@ -1118,7 +1242,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { _ => {} } if let Type::ClassDef(cls) = callee - && is_polars_series(cls) + && RuntimeClass::PolarsSeries.matches(cls) { return self .infer_series_dtype(arguments) @@ -1663,9 +1787,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return None; }; let class = underlying.class_object(); - if !is_polars_dataframe(class) && !is_polars_lazyframe(class) { - return None; - } + let kind = if is_polars_lazyframe(class) { + DataFrameKind::Polars + } else if RuntimeClass::NarwhalsLazyFrame.matches(class) { + DataFrameKind::Narwhals + } else { + eager_frame_kind(class)? + }; let (schema_cls, completeness) = match metadata { [Type::ClassDef(schema_cls)] => (schema_cls, SchemaCompleteness::Complete), [Type::ClassDef(schema_cls), tail] if tail.is_ellipsis_value() => { @@ -1678,7 +1806,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { underlying: underlying.clone(), columns: self.schema_class_columns(schema_cls)?, completeness, - kind: DataFrameKind::Polars, + kind, role: SchemaRole::Contract, } .to_type(), @@ -2070,7 +2198,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { fn infer_polars_concat( &self, arguments: &Arguments, - ) -> Option<(Vec<(Name, PolarsDType)>, SchemaCompleteness)> { + ) -> Option<(Vec<(Name, PolarsDType)>, SchemaCompleteness, DataFrameKind)> { let [items] = &arguments.args[..] else { return None; }; @@ -2083,25 +2211,30 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let schemas = elts .iter() .map(|e| match self.expr_infer(e, &self.error_swallower()) { - Type::DataFrame(schema) if schema.kind == DataFrameKind::Polars => { - Some((schema.columns, schema.completeness)) + Type::DataFrame(schema) if schema.kind.is_polars_api() => { + Some((schema.columns, schema.completeness, schema.kind)) } _ => None, }) .collect::>>()?; - let completeness = schemas - .iter() - .fold(SchemaCompleteness::Complete, |completeness, (_, next)| { - completeness.combine(*next) - }); + let completeness = schemas.iter().fold( + SchemaCompleteness::Complete, + |completeness, (_, next, _)| completeness.combine(*next), + ); let (first, rest) = schemas.split_first()?; + let kind = first.2; + // Concatenating frames from different libraries, e.g. a Polars frame with a Narwhals one, + // is not something either library supports, so fall back rather than merge across kinds. + if rest.iter().any(|(_, _, other)| *other != kind) { + return None; + } let columns = match how { ConcatHow::Vertical => rest .iter() - .all(|(columns, _)| columns == &first.0) + .all(|(columns, _, _)| columns == &first.0) .then(|| first.0.clone())?, ConcatHow::VerticalRelaxed => { - let names_match = rest.iter().all(|(columns, _)| { + let names_match = rest.iter().all(|(columns, _, _)| { columns.len() == first.0.len() && columns.iter().zip(&first.0).all(|((n, _), (m, _))| n == m) }); @@ -2113,15 +2246,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .iter() .enumerate() .map(|(i, (name, dtype))| { - let folded = rest.iter().try_fold(dtype.clone(), |acc, (columns, _)| { - acc.supertype(columns[i].1.clone()) - })?; + let folded = rest + .iter() + .try_fold(dtype.clone(), |acc, (columns, _, _)| { + acc.supertype(columns[i].1.clone()) + })?; Some((name.clone(), folded)) }) .collect::>>()? } }; - Some((columns, completeness)) + Some((columns, completeness, kind)) } /// Anchors on the first non-null element; only Polars reports later mismatches. @@ -2345,7 +2480,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } PolarsOutputSource::Expr(expr) => self .eval_polars_expr(expr, schema, errors) - .map(ExprValue::dtype), + .map(|value| value.column_dtype(schema.kind)), PolarsOutputSource::Unmodeled => None, }; if dtype.is_some() && !resolved_names.insert(output.name.clone()) { @@ -2407,7 +2542,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } if let [arg] = &positional_expressions(args)[..] @@ -2428,7 +2563,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { /// Remove statically named columns while preserving order. fn polars_drop(&self, base: &Type, args: &Arguments, errors: &ErrorCollector) -> Option { let schema = column_transform_schema(base, args)?; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } let positional = args @@ -2472,7 +2607,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let [Expr::Dict(mapping)] = &args.args[..] else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } let mut renames: SmallMap = @@ -2540,13 +2675,38 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } Some(PolarsExprMethod::Reducer(reducer)) => { let inner = self.eval_polars_expr(&attr.value, schema, errors)?; - return Some(ExprValue::Dtype(reducer.output_dtype(inner.dtype())?)); + return match reducer { + // `Identity` is `min`, `max`, `first`, or `last`. Narwhals + // guarantees that these keep the input dtype, so unlike the other + // reducers they are not limited to `Int64` and `Float64` inputs. + Reducer::Identity => { + Some(ExprValue::Dtype(inner.column_dtype(schema.kind))) + } + _ => { + // Narwhals does not guarantee the dtype of an integer sum, + // even of an `Int64` column. + let narwhals_unspecified = + matches!(reducer, Reducer::Sum) && inner.is_integer(); + computed_value( + schema.kind, + [inner], + narwhals_unspecified, + |[inner]| { + reducer + .output_dtype(inner.dtype()) + .map(ExprValue::Dtype) + }, + ) + } + }; } None => {} } } match self.polars_function(&call.func)? { - PolarsFunction::Len => Some(ExprValue::Dtype(PolarsDType::UInt32)), + PolarsFunction::Len => computed_value(schema.kind, [], false, |[]| { + Some(ExprValue::Dtype(PolarsDType::UInt32)) + }), PolarsFunction::Col => { let [arg] = &call.arguments.args[..] else { return None; @@ -2578,11 +2738,19 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { Expr::BinOp(binop) => { let a = self.eval_polars_expr(&binop.left, schema, errors)?; let b = self.eval_polars_expr(&binop.right, schema, errors)?; - combine_binop(binop.op, a, b) + // Narwhals does not guarantee the dtype of an integer power or bitwise operation. + let narwhals_unspecified = matches!( + binop.op, + Operator::Pow | Operator::BitAnd | Operator::BitOr | Operator::BitXor + ) && a.is_integer() + && b.is_integer(); + computed_value(schema.kind, [a, b], narwhals_unspecified, |[a, b]| { + combine_binop(binop.op, a, b) + }) } Expr::UnaryOp(unary) => { let a = self.eval_polars_expr(&unary.operand, schema, errors)?; - unary_value(unary.op, a) + computed_value(schema.kind, [a], false, |[a]| unary_value(unary.op, a)) } Expr::Compare(cmp) => { let ([op], [left, right]) = (&*cmp.ops, &*cmp.operands) else { @@ -2818,7 +2986,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { fn is_polars_expr_value(&self, expr: &Expr) -> bool { matches!( self.expr_infer(expr, &self.error_swallower()), - Type::ClassType(cls) if is_polars_expr(cls.class_object()) + Type::ClassType(cls) if is_polars_api_expr(cls.class_object()) ) } @@ -2832,7 +3000,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } match self.polars_projection(schema, args, errors)? { @@ -2873,7 +3041,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } // Validate every output name before emitting diagnostics. @@ -2948,7 +3116,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ColumnArg::Opaque => None, ColumnArg::Expr => Some( self.eval_polars_expr(expr, schema, errors) - .map_or(PolarsDType::Unknown, ExprValue::dtype), + .map_or(PolarsDType::Unknown, |value| { + value.column_dtype(schema.kind) + }), ), } } @@ -2971,7 +3141,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.expr_infer(expr, errors); let dtype = if self.polars_expr_aggregates(expr) { self.eval_polars_expr(expr, schema, errors) - .map_or(PolarsDType::Unknown, ExprValue::dtype) + .map_or(PolarsDType::Unknown, |value| { + value.column_dtype(schema.kind) + }) } else { // Evaluate for column-existence errors, but the list dtype is unmodeled. self.eval_polars_expr(expr, schema, errors); @@ -3091,7 +3263,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } for arg in args.args.iter() { @@ -3114,7 +3286,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } if let Some(ArgumentValue::Present(by)) = extract_argument(args, 0, "by") { @@ -3147,7 +3319,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } if let Some(ArgumentValue::Present(subset)) = extract_argument(args, 0, "subset") { @@ -3179,7 +3351,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } if let Some(ArgumentValue::Present(subset)) = extract_argument(args, 0, "subset") { @@ -3209,7 +3381,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } // A bare `Expr::Starred` is treated as a type form, so infer its value instead. @@ -3261,7 +3433,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars || !args.args.is_empty() { + if !schema.kind.is_polars_api() || !args.args.is_empty() { return None; } // Delegate keyword validation and the result class to the stub. @@ -3277,7 +3449,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ); match (conversion, result) { (PolarsFrameConversion::Lazy, Type::ClassType(cls)) - if is_polars_lazyframe(cls.class_object()) => + if is_lazy_frame(cls.class_object()) => { Some( DataFrameSchema { @@ -3289,7 +3461,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) } (PolarsFrameConversion::Collect, Type::ClassType(cls)) - if is_polars_dataframe(cls.class_object()) => + if eager_frame_kind(cls.class_object()).is_some() => { Some( DataFrameSchema { @@ -3310,7 +3482,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let [arg] = &args.args[..] else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } let columns = match arg { @@ -3357,7 +3529,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::DataFrame(schema) = base else { return None; }; - if schema.kind != DataFrameKind::Polars { + if !schema.kind.is_polars_api() { return None; } let [other_expr] = &args.args[..] else { @@ -3548,7 +3720,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { matches!( self.expr_infer(callee, &self.error_swallower()), Type::ClassDef(cls) - if is_polars_series(&cls) || RuntimeClass::PolarsDataFrameSeries.matches(&cls) + if is_polars_api_series(&cls) || RuntimeClass::PolarsDataFrameSeries.matches(&cls) ) } @@ -3624,7 +3796,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { None, ); match (dtype, result) { - (Some(dtype), Type::ClassType(cls)) if is_polars_series(cls.class_object()) => { + (Some(dtype), Type::ClassType(cls)) if is_polars_api_series(cls.class_object()) => { SeriesSchema { underlying: cls, dtype, diff --git a/pyrefly/lib/solver/subset.rs b/pyrefly/lib/solver/subset.rs index e203027570..7c21c50364 100644 --- a/pyrefly/lib/solver/subset.rs +++ b/pyrefly/lib/solver/subset.rs @@ -15,7 +15,6 @@ use itertools::Itertools; use itertools::izip; use pyrefly_python::dunder; use pyrefly_types::callable::Callable; -use pyrefly_types::data_frame::DataFrameKind; use pyrefly_types::dimension::Int; use pyrefly_types::dimension::ShapeError; use pyrefly_types::dimension::contains_var_in_type; @@ -2431,16 +2430,14 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> { &got_schema.underlying_type(), &want_schema.underlying_type(), )?; - if want_schema.kind == DataFrameKind::Polars && want_schema.is_contract() { + if want_schema.kind.is_polars_api() && want_schema.is_contract() { ok_or(got_schema.satisfies(want_schema), SubsetError::Other) } else { Ok(()) } } (Type::DataFrame(schema), _) => self.is_subset_eq(&schema.underlying_type(), want), - (_, Type::DataFrame(schema)) - if schema.kind == DataFrameKind::Polars && schema.is_contract() => - { + (_, Type::DataFrame(schema)) if schema.kind.is_polars_api() && schema.is_contract() => { Err(SubsetError::Other) } (_, Type::DataFrame(schema)) => self.is_subset_eq(got, &schema.underlying_type()), diff --git a/pyrefly/lib/state/lsp/dict_completions.rs b/pyrefly/lib/state/lsp/dict_completions.rs index 6f979d43c3..45baa17156 100644 --- a/pyrefly/lib/state/lsp/dict_completions.rs +++ b/pyrefly/lib/state/lsp/dict_completions.rs @@ -286,28 +286,19 @@ impl<'a> Transaction<'a> { inside_column_helper: bool, ) -> Option { match ty { - Type::DataFrame(schema) => Some(match (schema.kind, method, slot) { - (DataFrameKind::Polars, "select" | "with_columns", _) => true, - (DataFrameKind::Polars, "drop" | "filter", ArgumentSlot::Positional) => true, - ( - DataFrameKind::Polars, - "filter", - ArgumentSlot::Keyword(_) | ArgumentSlot::UnpackedKeyword, - ) => inside_column_helper, - ( - DataFrameKind::Polars, - "sort", - ArgumentSlot::Positional | ArgumentSlot::Keyword("by"), - ) => true, - ( - DataFrameKind::Polars, - "agg", - ArgumentSlot::Positional | ArgumentSlot::Keyword(_), - ) => true, - (DataFrameKind::Polars, "group_by" | "groupby", ArgumentSlot::Positional) => true, - (DataFrameKind::Polars, "group_by" | "groupby", ArgumentSlot::Keyword(name)) => { - name != "maintain_order" + Type::DataFrame(schema) if schema.kind.is_polars_api() => Some(match (method, slot) { + ("select" | "with_columns", _) => true, + ("drop" | "filter", ArgumentSlot::Positional) => true, + ("filter", ArgumentSlot::Keyword(_) | ArgumentSlot::UnpackedKeyword) => { + inside_column_helper } + ("sort", ArgumentSlot::Positional | ArgumentSlot::Keyword("by")) => true, + ("agg", ArgumentSlot::Positional | ArgumentSlot::Keyword(_)) => true, + ("group_by" | "groupby", ArgumentSlot::Positional) => true, + ("group_by" | "groupby", ArgumentSlot::Keyword(name)) => name != "maintain_order", + _ => false, + }), + Type::DataFrame(schema) => Some(match (schema.kind, method, slot) { ( DataFrameKind::Pandas, "drop", diff --git a/pyrefly/lib/test/mod.rs b/pyrefly/lib/test/mod.rs index f25778ea0a..8fa5c2cfaf 100644 --- a/pyrefly/lib/test/mod.rs +++ b/pyrefly/lib/test/mod.rs @@ -58,6 +58,7 @@ mod mro; mod named_ints; mod named_tuple; mod narrow; +mod narwhals; mod natural; mod new_type; mod nn_module; diff --git a/pyrefly/lib/test/narwhals.rs b/pyrefly/lib/test/narwhals.rs new file mode 100644 index 0000000000..4801532e17 --- /dev/null +++ b/pyrefly/lib/test/narwhals.rs @@ -0,0 +1,10 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#![cfg(test)] +mod dataframe; +mod util; diff --git a/pyrefly/lib/test/narwhals/dataframe.rs b/pyrefly/lib/test/narwhals/dataframe.rs new file mode 100644 index 0000000000..1aa42ad123 --- /dev/null +++ b/pyrefly/lib/test/narwhals/dataframe.rs @@ -0,0 +1,306 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use crate::narwhals_testcase; +use crate::test::narwhals::util::env_with_narwhals_and_polars_stubs; +use crate::testcase; + +narwhals_testcase!( + test_select_narrows_columns, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.select("a")) # E: revealed type: DataFrame[a: Int64] + df.select("missing") # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_drop_and_rename, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.drop("b")) # E: revealed type: DataFrame[a: Int64] + reveal_type(df.rename({"a": "c"})) # E: revealed type: DataFrame[c: Int64, b: String] +"#, +); + +narwhals_testcase!( + test_with_columns_adds_column, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.with_columns(b=nw.col("a") * nw.lit(2))) # E: revealed type: DataFrame[a: Int64, b: Int64] + df.with_columns(c=nw.col("missing")) # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_row_transforms_preserve_schema, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.filter(nw.col("a") > 0)) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.sort("a")) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.head(3)) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.drop_nulls()) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.unique()) # E: revealed type: DataFrame[a: Int64, b: String] +"#, +); + +narwhals_testcase!( + test_lazy_collect_round_trips_schema, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.lazy()) # E: revealed type: LazyFrame[a: Int64, b: String] + reveal_type(df.lazy().select("a")) # E: revealed type: LazyFrame[a: Int64] + reveal_type(df.lazy().collect()) # E: revealed type: DataFrame[a: Int64, b: String] +"#, +); + +narwhals_testcase!( + test_column_access_and_unknown_column, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df["a"]) # E: revealed type: Series[Int64] + reveal_type(df.get_column("b")) # E: revealed type: Series[String] + df["missing"] # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_join_merges_schemas, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class LeftSchema: + key: nw.Int64 + a: nw.String +class RightSchema: + key: nw.Int64 + b: nw.Float64 +def f( + left: Annotated[nw.DataFrame, LeftSchema], right: Annotated[nw.DataFrame, RightSchema] +) -> None: + reveal_type(left.join(right, on="key")) # E: revealed type: DataFrame[key: Int64, a: String, b: Float64] +"#, +); + +narwhals_testcase!( + test_group_by_agg, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + g: nw.String + v: nw.Int64 +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.group_by("g").agg(nw.col("v").mean())) # E: revealed type: DataFrame[g: String, v: Float64] +"#, +); + +narwhals_testcase!( + test_cast_updates_dtypes, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.cast({"a": nw.Float64})) # E: revealed type: DataFrame[a: Float64, b: String] + reveal_type(df.lazy().cast(nw.String)) # E: revealed type: LazyFrame[a: String, b: String] + df.cast({"missing": nw.Float64}) # E: Column `missing` is not in the DataFrame schema +"#, +); + +// Narwhals guarantees computed numeric dtypes only for `Int64` and `Float64`, so other computed +// numeric dtypes are unknown. Comparisons, casts, and `min`/`max` keep their dtype. +narwhals_testcase!( + test_unspecified_dtypes_are_unknown, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + g: nw.String + small: nw.Int8 + big: nw.Int64 + flag: nw.Boolean +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.select(a=nw.col("small") + 1, b=nw.col("big") + 1, c=nw.lit(1))) # E: revealed type: DataFrame[a: Unknown, b: Int64, c: Unknown] + reveal_type(df.select(a=nw.col("big") * 2, b=nw.col("big") ** 2, c=nw.col("big") & nw.col("big"))) # E: revealed type: DataFrame[a: Int64, b: Unknown, c: Unknown] + reveal_type(df.select(a=nw.col("small") > 0, b=(nw.col("small") + 1).cast(nw.Int64))) # E: revealed type: DataFrame[a: Boolean, b: Int64] + reveal_type(df.group_by("g").agg(a=nw.col("small").sum(), b=nw.col("big").sum(), c=nw.col("flag").sum())) # E: revealed type: DataFrame[g: String, a: Unknown, b: Unknown, c: Unknown] + reveal_type(df.group_by("g").agg(a=nw.col("big").count(), b=nw.col("small").max(), c=nw.col("big").mean())) # E: revealed type: DataFrame[g: String, a: Unknown, b: Int8, c: Float64] +"#, +); + +// The dtype of `len` is unknown because Narwhals does not guarantee it. +// `narwhals.len` is `narwhals.functions.len_` imported under its public name, since `len_` +// avoids shadowing the builtin; the stable namespaces instead define their own `len` functions. +narwhals_testcase!( + test_len_resolves_from_top_level_and_stable_namespaces, + r#" +import narwhals as nw +import narwhals.stable.v1 as nw_v1 +import narwhals.stable.v2 as nw_v2 +from typing import Annotated, reveal_type +class MySchema: + g: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df.group_by("g").agg(nw.len())) # E: revealed type: DataFrame[g: String, len: Unknown] + +class MySchemaV1: + g: nw_v1.String +def f_v1(df: Annotated[nw_v1.DataFrame, MySchemaV1]) -> None: + reveal_type(df.group_by("g").agg(nw_v1.len())) # E: revealed type: DataFrame[g: String, len: Unknown] + +class MySchemaV2: + g: nw_v2.String +def f_v2(df: Annotated[nw_v2.DataFrame, MySchemaV2]) -> None: + reveal_type(df.group_by("g").agg(nw_v2.len())) # E: revealed type: DataFrame[g: String, len: Unknown] +"#, +); + +narwhals_testcase!( + test_lazy_group_by_agg, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + g: nw.String + v: nw.Int64 +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + lazy = df.lazy() + reveal_type(lazy.group_by("g").agg(nw.col("v").mean())) # E: revealed type: LazyFrame[g: String, v: Float64] + lazy.group_by("missing").agg(nw.col("v").sum()) # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_duplicate_output_column_is_rejected, + r#" +import narwhals as nw +from typing import Annotated +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + df.select("a", "a") # E: Operation produces duplicate column `a` +"#, +); + +narwhals_testcase!( + test_concat_preserves_schema, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 +def f(a: Annotated[nw.DataFrame, MySchema], b: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(nw.concat([a, b])) # E: revealed type: DataFrame[a: Int64] +"#, +); + +narwhals_testcase!( + test_stable_v1_namespace_tracks_schema, + r#" +import narwhals.stable.v1 as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.select("a")) # E: revealed type: DataFrame[a: Int64] + df.select("missing") # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_stable_v2_namespace_tracks_schema, + r#" +import narwhals.stable.v2 as nw +from typing import Annotated, reveal_type +class MySchema: + a: nw.Int64 + b: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df) # E: revealed type: DataFrame[a: Int64, b: String] + reveal_type(df.select("a")) # E: revealed type: DataFrame[a: Int64] + df.select("missing") # E: Column `missing` is not in the DataFrame schema +"#, +); + +narwhals_testcase!( + test_dataframe_schema_annotation, + r#" +import narwhals as nw +from typing import Annotated, reveal_type +class MySchema: + price: nw.Float64 + asset: nw.String +def f(df: Annotated[nw.DataFrame, MySchema]) -> None: + reveal_type(df) # E: revealed type: DataFrame[price: Float64, asset: String] + reveal_type(df["price"]) # E: revealed type: Series[Float64] + df["missing"] # E: Column `missing` is not in the DataFrame schema +"#, +); + +// Schema contracts are per-library: a Polars frame does not satisfy a Narwhals schema +// declaration, even when its inferred columns match, because the two libraries model +// distinct runtime frame types. +testcase!( + test_schema_contracts_do_not_cross_libraries, + env_with_narwhals_and_polars_stubs(), + r#" +import narwhals as nw +import polars as pl +from typing import Annotated +class MySchema: + a: nw.Int64 +def takes_narwhals(df: Annotated[nw.DataFrame, MySchema]) -> None: ... +def also_takes_narwhals(df: Annotated[nw.DataFrame, MySchema]) -> None: + takes_narwhals(df) +takes_narwhals(pl.DataFrame({"a": [1]})) # E: `polars.dataframe.frame.DataFrame[a: Int64]` is not assignable to parameter `df` with type `narwhals.dataframe.DataFrame[a: Int64]` +"#, +); + +narwhals_testcase!( + test_opaque_frame_falls_back, + r#" +import narwhals as nw +from typing import reveal_type +def f(df: nw.DataFrame) -> None: + reveal_type(df.select("anything")) # E: revealed type: DataFrame +"#, +); diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/__init__.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/__init__.pyi new file mode 100644 index 0000000000..4f640b9c1c --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/__init__.pyi @@ -0,0 +1,30 @@ +"""Minimal Narwhals stubs for schema-tracking tests.""" + +from narwhals.dataframe import DataFrame as DataFrame, LazyFrame as LazyFrame +from narwhals.dtypes import ( + Array as Array, + Boolean as Boolean, + Float32 as Float32, + Float64 as Float64, + Int8 as Int8, + Int16 as Int16, + Int32 as Int32, + Int64 as Int64, + Int128 as Int128, + List as List, + String as String, + Struct as Struct, + UInt8 as UInt8, + UInt64 as UInt64, + UInt128 as UInt128, +) +from narwhals.expr import Expr as Expr +from narwhals.functions import ( + col as col, + concat as concat, + len_ as len, + lit as lit, + when as when, +) +from narwhals.schema import Schema as Schema +from narwhals.series import Series as Series diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/dataframe.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/dataframe.pyi new file mode 100644 index 0000000000..7e84b2e901 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/dataframe.pyi @@ -0,0 +1,75 @@ +from typing import Iterator, Self + +from narwhals.series import Series + +class DataFrame: + columns: list[str] + + def __getitem__(self, key: str) -> Series: ... + def __iter__(self) -> Iterator[Series]: ... + def __contains__(self, key: str) -> bool: ... + def get_column(self, name: str) -> Series: ... + def head(self, n: int = 5) -> Self: ... + def tail(self, n: int = 5) -> Self: ... + def select(self, *exprs: object, **named_exprs: object) -> Self: ... + def drop(self, *columns: object, strict: bool = True) -> Self: ... + def rename(self, mapping: object) -> Self: ... + def cast(self, dtypes: object) -> Self: ... + def with_columns(self, *exprs: object, **named_exprs: object) -> Self: ... + def filter(self, *predicates: object, **constraints: object) -> Self: ... + def sort(self, by: object, *more: object, descending: bool = False) -> Self: ... + def unique( + self, + subset: object = None, + *, + keep: str = "any", + maintain_order: bool | None = None, + ) -> Self: ... + def drop_nulls(self, subset: object = None) -> Self: ... + def join( + self, + other: Self, + on: object = None, + how: str = "inner", + *, + left_on: object = None, + right_on: object = None, + suffix: str = "_right", + ) -> Self: ... + def group_by(self, *keys: object, drop_null_keys: bool = False) -> GroupBy: ... + def lazy(self) -> LazyFrame: ... + def to_native(self) -> object: ... + +class LazyFrame: + columns: list[str] + + def select(self, *exprs: object, **named_exprs: object) -> Self: ... + def drop(self, *columns: object, strict: bool = True) -> Self: ... + def rename(self, mapping: object) -> Self: ... + def cast(self, dtypes: object) -> Self: ... + def with_columns(self, *exprs: object, **named_exprs: object) -> Self: ... + def filter(self, *predicates: object, **constraints: object) -> Self: ... + def sort(self, by: object, *more: object, descending: bool = False) -> Self: ... + def unique(self, subset: object = None, *, keep: str = "any") -> Self: ... + def drop_nulls(self, subset: object = None) -> Self: ... + def join( + self, + other: Self, + on: object = None, + how: str = "inner", + *, + left_on: object = None, + right_on: object = None, + suffix: str = "_right", + ) -> Self: ... + def group_by(self, *keys: object, drop_null_keys: bool = False) -> LazyGroupBy: ... + def head(self, n: int = 5) -> Self: ... + def lazy(self) -> Self: ... + def collect(self) -> DataFrame: ... + def to_native(self) -> object: ... + +class GroupBy: + def agg(self, *aggs: object, **named_aggs: object) -> DataFrame: ... + +class LazyGroupBy: + def agg(self, *aggs: object, **named_aggs: object) -> LazyFrame: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/dtypes.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/dtypes.pyi new file mode 100644 index 0000000000..b5a8a1b15e --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/dtypes.pyi @@ -0,0 +1,23 @@ +class Int8: ... +class Int16: ... +class Int32: ... +class Int64: ... +class Int128: ... +class UInt8: ... +class UInt64: ... +class UInt128: ... +class Float32: ... +class Float64: ... +class String: ... +class Boolean: ... + +class Array: + def __init__( + self, inner: object, shape: int | tuple[int, ...] | None = None + ) -> None: ... + +class List: + def __init__(self, inner: object) -> None: ... + +class Struct: + def __init__(self, fields: dict[str, object]) -> None: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/expr.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/expr.pyi new file mode 100644 index 0000000000..9f9a2c1248 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/expr.pyi @@ -0,0 +1,34 @@ +from typing import Any + +class Expr: + def __add__(self, other: Any) -> Expr: ... + def __radd__(self, other: Any) -> Expr: ... + def __sub__(self, other: Any) -> Expr: ... + def __mul__(self, other: Any) -> Expr: ... + def __truediv__(self, other: Any) -> Expr: ... + def __floordiv__(self, other: Any) -> Expr: ... + def __mod__(self, other: Any) -> Expr: ... + def __pow__(self, other: Any) -> Expr: ... + def __rpow__(self, other: Any) -> Expr: ... + def __and__(self, other: Any) -> Expr: ... + def __or__(self, other: Any) -> Expr: ... + def __xor__(self, other: Any) -> Expr: ... + def __neg__(self) -> Expr: ... + def __invert__(self) -> Expr: ... + def __gt__(self, other: Any) -> Expr: ... + def __ge__(self, other: Any) -> Expr: ... + def __lt__(self, other: Any) -> Expr: ... + def __le__(self, other: Any) -> Expr: ... + def alias(self, name: str) -> Expr: ... + def cast(self, dtype: Any) -> Expr: ... + def sum(self) -> Expr: ... + def mean(self) -> Expr: ... + def median(self) -> Expr: ... + def std(self, *, ddof: int = 1) -> Expr: ... + def var(self, *, ddof: int = 1) -> Expr: ... + def min(self) -> Expr: ... + def max(self) -> Expr: ... + def first(self) -> Expr: ... + def last(self) -> Expr: ... + def count(self) -> Expr: ... + def n_unique(self) -> Expr: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/functions.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/functions.pyi new file mode 100644 index 0000000000..5da6233799 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/functions.pyi @@ -0,0 +1,17 @@ +from typing import Iterable + +from narwhals.dataframe import DataFrame +from narwhals.expr import Expr + +def col(*names: str | Iterable[str]) -> Expr: ... +def lit(value: object, dtype: object = None) -> Expr: ... +def len_() -> Expr: ... +def concat(items: Iterable[DataFrame], *, how: str = "vertical") -> DataFrame: ... + +class When: + def then(self, value: object) -> Then: ... + +class Then(Expr): + def otherwise(self, value: object) -> Expr: ... + +def when(*predicates: object) -> When: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/schema.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/schema.pyi new file mode 100644 index 0000000000..69028881f4 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/schema.pyi @@ -0,0 +1,2 @@ +class Schema: + def __init__(self, schema: object = None) -> None: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/series.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/series.pyi new file mode 100644 index 0000000000..2e8482c26e --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/series.pyi @@ -0,0 +1,9 @@ +from typing import Any, overload + +class Series: + @overload + def __getitem__(self, key: int) -> Any: ... + @overload + def __getitem__(self, key: slice) -> Series: ... + def __or__(self, other: Series) -> Series: ... + def to_native(self) -> object: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/stable/__init__.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/stable/__init__.pyi new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v1/__init__.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v1/__init__.pyi new file mode 100644 index 0000000000..902cecbc46 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v1/__init__.pyi @@ -0,0 +1,38 @@ +"""The stable namespace re-declares the public API as subclasses of the main one.""" + +from typing import Iterable + +from narwhals.dataframe import DataFrame as NwDataFrame, LazyFrame as NwLazyFrame +from narwhals.dtypes import ( + Array as Array, + Boolean as Boolean, + Float32 as Float32, + Float64 as Float64, + Int8 as Int8, + Int16 as Int16, + Int32 as Int32, + Int64 as Int64, + Int128 as Int128, + List as List, + String as String, + Struct as Struct, + UInt8 as UInt8, + UInt64 as UInt64, + UInt128 as UInt128, +) +from narwhals.expr import Expr as NwExpr +from narwhals.functions import Then, When +from narwhals.schema import Schema as NwSchema +from narwhals.series import Series as NwSeries + +class DataFrame(NwDataFrame): ... +class LazyFrame(NwLazyFrame): ... +class Series(NwSeries): ... +class Expr(NwExpr): ... +class Schema(NwSchema): ... + +def col(*names: str | Iterable[str]) -> Expr: ... +def lit(value: object, dtype: object = None) -> Expr: ... +def len() -> Expr: ... +def concat(items: Iterable[DataFrame], *, how: str = "vertical") -> DataFrame: ... +def when(*predicates: object) -> When: ... diff --git a/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v2/__init__.pyi b/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v2/__init__.pyi new file mode 100644 index 0000000000..902cecbc46 --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/stable/v2/__init__.pyi @@ -0,0 +1,38 @@ +"""The stable namespace re-declares the public API as subclasses of the main one.""" + +from typing import Iterable + +from narwhals.dataframe import DataFrame as NwDataFrame, LazyFrame as NwLazyFrame +from narwhals.dtypes import ( + Array as Array, + Boolean as Boolean, + Float32 as Float32, + Float64 as Float64, + Int8 as Int8, + Int16 as Int16, + Int32 as Int32, + Int64 as Int64, + Int128 as Int128, + List as List, + String as String, + Struct as Struct, + UInt8 as UInt8, + UInt64 as UInt64, + UInt128 as UInt128, +) +from narwhals.expr import Expr as NwExpr +from narwhals.functions import Then, When +from narwhals.schema import Schema as NwSchema +from narwhals.series import Series as NwSeries + +class DataFrame(NwDataFrame): ... +class LazyFrame(NwLazyFrame): ... +class Series(NwSeries): ... +class Expr(NwExpr): ... +class Schema(NwSchema): ... + +def col(*names: str | Iterable[str]) -> Expr: ... +def lit(value: object, dtype: object = None) -> Expr: ... +def len() -> Expr: ... +def concat(items: Iterable[DataFrame], *, how: str = "vertical") -> DataFrame: ... +def when(*predicates: object) -> When: ... diff --git a/pyrefly/lib/test/narwhals/util.rs b/pyrefly/lib/test/narwhals/util.rs new file mode 100644 index 0000000000..6b1ec7e34b --- /dev/null +++ b/pyrefly/lib/test/narwhals/util.rs @@ -0,0 +1,49 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use std::env; + +use crate::test::util::TestEnv; + +/// Returns an environment containing minimal stubs with the real Narwhals qualified names. +pub fn env_with_narwhals_stubs() -> TestEnv { + let path = env::var("NARWHALS_TEST_PATH").expect("NARWHALS_TEST_PATH must be set"); + TestEnv::new_with_site_package_paths(&[&path]) +} + +/// Narwhals wraps a backend frame, so some behavior is only observable with a backend present. +pub fn env_with_narwhals_and_polars_stubs() -> TestEnv { + let narwhals = env::var("NARWHALS_TEST_PATH").expect("NARWHALS_TEST_PATH must be set"); + let polars = env::var("POLARS_TEST_PATH").expect("POLARS_TEST_PATH must be set"); + TestEnv::new_with_site_package_paths(&[&narwhals, &polars]) +} + +#[macro_export] +macro_rules! narwhals_testcase { + (bug = $explanation:literal, $name:ident, $contents:literal,) => { + #[test] + fn $name() -> anyhow::Result<()> { + $crate::test::util::testcase_for_macro( + $crate::test::narwhals::util::env_with_narwhals_stubs(), + $contents, + file!(), + line!(), + ) + } + }; + ($name:ident, $contents:literal,) => { + #[test] + fn $name() -> anyhow::Result<()> { + $crate::test::util::testcase_for_macro( + $crate::test::narwhals::util::env_with_narwhals_stubs(), + $contents, + file!(), + line!() - 1, + ) + } + }; +} diff --git a/website/docs/dataframes.mdx b/website/docs/dataframes.mdx index c1c1b7abb0..04c07b938f 100644 --- a/website/docs/dataframes.mdx +++ b/website/docs/dataframes.mdx @@ -1,6 +1,6 @@ --- title: DataFrames -description: Pyrefly support for Polars and pandas DataFrames, including column and dtype inference. +description: Pyrefly support for Polars, Narwhals, and pandas DataFrames, including column and dtype inference. --- {/* @@ -34,8 +34,9 @@ Inline comments show how Pyrefly can infer additional information for DataFrames and Series. Pyrefly includes **built-in support** for DataFrames from -[Polars](https://docs.pola.rs/) and -[pandas](https://pandas.pydata.org/docs/), two popular Python libraries for +[Polars](https://docs.pola.rs/), +[Narwhals](https://narwhals-dev.github.io/narwhals/), and +[pandas](https://pandas.pydata.org/docs/), three popular Python libraries for working with tabular data. Pyrefly tracks column names and data types (dtypes), reports accesses to columns that cannot exist, and follows schema changes through common operations. @@ -47,7 +48,7 @@ through common operations. DataFrame inference works automatically in Pyrefly, without any special configuration. -1. Install `polars` or `pandas` in your Python environment. +1. Install `polars`, `narwhals`, or `pandas` in your Python environment. 2. Install `pyrefly`. 3. Write DataFrame code as usual. 4. Run Pyrefly or use the Pyrefly language server in your editor. @@ -121,10 +122,11 @@ Pyrefly infers a dtype for every known column by following the library's widening and coercion rules. When no dtype can be determined, the column keeps its name and takes the `Unknown` dtype, so an access returns `Series[Unknown]`. -### How Polars and pandas Are Treated Differently +### How the Libraries Are Treated Differently -The two libraries get different levels of precision because they offer different -mutation guarantees. +The libraries get different levels of precision because they offer different +mutation guarantees. Narwhals exposes a subset of the Polars API with the same +semantics, so Pyrefly treats Narwhals frames like Polars frames. In Polars, the common transformations such as `select`, `drop`, `rename`, and `with_columns` are immutable and return a new DataFrame with a transformed @@ -147,10 +149,11 @@ one. ## Supported Features -For brevity, the examples below omit the imports for `polars`, `pandas`, and the -`typing` names they use, such as `reveal_type` and `Literal`. Unless a section -says otherwise, these features are Polars only. pandas support is described in -[pandas Support](#pandas-support). +For brevity, the examples below omit the imports for `polars`, `narwhals`, +`pandas`, and the `typing` names they use, such as `reveal_type` and `Literal`. +Unless a section says otherwise, these features are Polars only. Narwhals +support is described in [Narwhals Support](#narwhals-support), and pandas +support in [pandas Support](#pandas-support). ### DataFrame Construction @@ -550,11 +553,62 @@ df.insert_column(1, pl.Series("active", [True])) reveal_type(df) # DataFrame[name: String, active: Unknown, age: Int64] ``` +### Narwhals Support + +[Narwhals](https://narwhals-dev.github.io/narwhals/) is a compatibility layer +that exposes a subset of the Polars API over several backends, with the same +method semantics and dtypes. The features above also apply to Narwhals frames, +for the methods Narwhals has. Pyrefly treats frames from `narwhals.stable.v1` +and `narwhals.stable.v2` the same way as frames from `narwhals`. + +Narwhals has no public frame constructor, and Pyrefly does not model +`from_native`, `from_dict`, or the Narwhals IO functions, so those calls produce +an opaque frame. A Narwhals frame gets a schema only from a +[schema declaration](#schema-declaration) on a parameter or return type. + +```python +class Sales: + region: nw.String + units: nw.Int64 + +def process(df: Annotated[nw.DataFrame, Sales]) -> None: + reveal_type(df) # DataFrame[region: String, units: Int64] + reveal_type(df.select("units")) # DataFrame[units: Int64] + df["missing"] # Error unknown-column +``` + +A Polars frame does not satisfy a Narwhals schema declaration, and a Narwhals +frame does not satisfy a Polars one. Methods that Narwhals lacks, such as +frame-level `fill_null` and the in-place column APIs, are supported only on +Polars frames. + +Narwhals does not guarantee the dtype of every computed column, so Pyrefly +infers the dtype of an arithmetic expression or aggregation only when every +numeric input and the result are `Int64` or `Float64`. Otherwise the column has +the `Unknown` dtype. Integer sums, powers, bitwise operations, and `nw.len()` +are `Unknown` even on `Int64` columns. Comparisons, casts, `min`, +`max`, `first`, and `last` keep their dtype. + +```python +class Sales: + region: nw.String + units: nw.Int64 + discount: nw.Float32 + +def summarize(df: Annotated[nw.DataFrame, Sales]) -> None: + out = df.group_by("region").agg( + nw.col("units").mean(), + nw.col("discount").mean(), + nw.col("units").sum().alias("total"), + nw.len(), + ) + reveal_type(out) # DataFrame[region: String, units: Float64, discount: Unknown, total: Unknown, len: Unknown] +``` + ### pandas Support -Everything above this section describes Polars. pandas gets a smaller feature -set, because its open-ended mutation model means an inferred schema is always -partial. +pandas gets a smaller feature set than Polars and Narwhals, because its +open-ended mutation model means an inferred schema is always partial. Pyrefly infers visible primitive columns from a pandas dictionary constructor passed positionally or through `data=`. A static `columns=` list projects and @@ -587,8 +641,8 @@ configuration. - [`unknown-column`](../error-kinds/#unknown-column) reports a statically known name that is absent from a complete schema. - [`column-type-mismatch`](../error-kinds/#column-type-mismatch) reports a - strict Polars construction value that does not fit the dtype established by - the first non-null value. + strict Polars or Narwhals construction value that does not fit the dtype + established by the first non-null value. - [`column-schema-mismatch`](../error-kinds/#column-schema-mismatch) reports statically known data names that do not match declared schema names. - [`duplicate-column`](../error-kinds/#duplicate-column) reports an operation