WIP: Narwhals support

Author: MarcoGorelli

.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 }

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

pyrefly/lib/alt/expr.rs

@@ -18,7 +18,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;
@@ -88,7 +87,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::shape_extension::is_int_tuple_bound;
@@ -3833,7 +3832,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);
@@ -3849,7 +3848,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(_) => {}
@@ -3876,7 +3875,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,