From 9aef418f5e7b8d2dd16c09c627e4d08524364d9d 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] WIP: 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 | 197 ++++++++++---- pyrefly/lib/solver/subset.rs | 7 +- pyrefly/lib/state/lsp/dict_completions.rs | 27 +- pyrefly/lib/test/mod.rs | 1 + pyrefly/lib/test/narwhals.rs | 10 + pyrefly/lib/test/narwhals/dataframe.rs | 248 ++++++++++++++++++ .../third-party/narwhals/__init__.pyi | 30 +++ .../third-party/narwhals/dataframe.pyi | 73 ++++++ .../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 +++ pyrefly/lib/test/narwhals/util.rs | 49 ++++ website/docs/dataframes.mdx | 64 +++-- 20 files changed, 746 insertions(+), 100 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/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 25c19646fe..a66477311d 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 a9cff8b3dd..b12790dabd 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; @@ -89,7 +88,7 @@ use crate::alt::answers_solver::TypeCheckOptions; 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::regular_nested_list::regular_nested_list; @@ -3806,7 +3805,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); @@ -3822,7 +3821,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(_) => {} @@ -3849,7 +3848,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 fd51a5519b..b1b59adcb5 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()) @@ -321,7 +374,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, @@ -739,7 +815,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> { @@ -1096,9 +1172,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, }, ); @@ -1106,7 +1182,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) @@ -1649,9 +1725,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let Type::ClassType(underlying) = inner else { return None; }; - if !is_polars_dataframe(underlying.class_object()) { - return None; - } + let kind = eager_frame_kind(underlying.class_object())?; let (schema_cls, completeness) = match metadata { [Type::ClassDef(schema_cls)] => (schema_cls, SchemaCompleteness::Complete), [Type::ClassDef(schema_cls), tail] if tail.is_ellipsis_value() => { @@ -1664,7 +1738,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(), @@ -2056,7 +2130,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; }; @@ -2069,25 +2143,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) }); @@ -2099,15 +2178,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. @@ -2393,7 +2474,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)[..] @@ -2414,7 +2495,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 @@ -2458,7 +2539,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 = @@ -2804,7 +2885,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()) ) } @@ -2818,7 +2899,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)? { @@ -2859,7 +2940,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. @@ -3078,7 +3159,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. @@ -3130,7 +3211,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. @@ -3146,7 +3227,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 { @@ -3158,7 +3239,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 { @@ -3226,7 +3307,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 { @@ -3417,7 +3498,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) ) } @@ -3493,7 +3574,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 1570ba0b39..594aa1155d 100644 --- a/pyrefly/lib/solver/subset.rs +++ b/pyrefly/lib/solver/subset.rs @@ -16,7 +16,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; @@ -2306,16 +2305,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 f0cc16f3e7..ffcb21d441 100644 --- a/pyrefly/lib/state/lsp/dict_completions.rs +++ b/pyrefly/lib/state/lsp/dict_completions.rs @@ -286,23 +286,18 @@ 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, "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, + ("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 ceea84c435..cc6c6dd21e 100644 --- a/pyrefly/lib/test/mod.rs +++ b/pyrefly/lib/test/mod.rs @@ -54,6 +54,7 @@ mod marshmallow; mod mro; 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..40aa871327 --- /dev/null +++ b/pyrefly/lib/test/narwhals/dataframe.rs @@ -0,0 +1,248 @@ +/* + * 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").sum())) # E: revealed type: DataFrame[g: String, v: Int64] +"#, +); + +// `narwhals.len` is `narwhals.functions.len_` imported under its public name, since `len_` +// avoids shadowing the builtin; `narwhals.stable.v1.len` is instead its own function. +narwhals_testcase!( + test_len_resolves_from_top_level_and_stable_namespaces, + r#" +import narwhals as nw +import narwhals.stable.v1 as nw_v1 +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: UInt32] + +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: UInt32] +"#, +); + +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").sum())) # E: revealed type: LazyFrame[g: String, v: Int64] + 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_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..2968f10ccc --- /dev/null +++ b/pyrefly/lib/test/narwhals/third-party/narwhals/dataframe.pyi @@ -0,0 +1,73 @@ +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 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 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/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..2c8efce733 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,12 @@ 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 wraps a backend frame behind a subset of the +Polars API with the same semantics, so it is treated like Polars throughout, +and the rest of this section contrasts Polars with pandas. 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 +150,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 +554,39 @@ 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. Because that +subset keeps Polars' method semantics and dtype system, every feature above +applies to Narwhals frames too, for the operations Narwhals offers. Frames +reached through `narwhals.stable.v1` and `narwhals.stable.v2` are tracked the +same way as those from the top-level `narwhals` namespace. + +Narwhals has no public frame constructor, and `from_native`, `from_dict`, and +the rest of the Narwhals IO surface are not modeled, so those calls produce an +opaque frame. A tracked schema starts 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 +``` + +Schemas do not cross libraries: a Polars frame does not satisfy a Narwhals +schema declaration, or the reverse. Operations Narwhals does not have, such as +frame-level `cast` and `fill_null` or the in-place column APIs, are Polars only. + ### 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 +619,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