From 81938efd3752744a4404f558800a1ec384fc6b26 Mon Sep 17 00:00:00 2001 From: FBruzzesi <francesco.bruzzesi.93@gmail.com> Date: Mon, 27 Jan 2025 16:47:04 +0100 Subject: [PATCH] not much but a honest job --- narwhals/_arrow/expr_name.py | 101 +++++++-------------------- narwhals/_dask/expr_name.py | 71 +++++-------------- narwhals/_duckdb/expr_name.py | 66 +++++------------- narwhals/_pandas_like/expr_name.py | 108 ++++++++--------------------- narwhals/_spark_like/expr_name.py | 64 +++++------------ 5 files changed, 109 insertions(+), 301 deletions(-) diff --git a/narwhals/_arrow/expr_name.py b/narwhals/_arrow/expr_name.py index a8741bf61..48cb77302 100644 --- a/narwhals/_arrow/expr_name.py +++ b/narwhals/_arrow/expr_name.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING from typing import Callable +from typing import Sequence if TYPE_CHECKING: from typing_extensions import Self @@ -14,107 +15,59 @@ def __init__(self: Self, expr: ArrowExpr) -> None: self._compliant_expr = expr def keep(self: Self) -> ArrowExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(name) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: name, alias_output_names=None, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, ) def map(self: Self, function: Callable[[str], str]) -> ArrowExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(function(str(name))) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: function(str(name)), alias_output_names=lambda output_names: [ function(str(name)) for name in output_names ], - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "function": function}, ) def prefix(self: Self, prefix: str) -> ArrowExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(f"{prefix}{name}") - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: f"{prefix}{name}", alias_output_names=lambda output_names: [ f"{prefix}{output_name}" for output_name in output_names ], - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "prefix": prefix}, ) def suffix(self: Self, suffix: str) -> ArrowExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(f"{name}{suffix}") - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: f"{name}{suffix}", alias_output_names=lambda output_names: [ f"{output_name}{suffix}" for output_name in output_names ], - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "suffix": suffix}, ) def to_lowercase(self: Self) -> ArrowExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(str(name).lower()) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: str(name).lower(), alias_output_names=lambda output_names: [ str(name).lower() for name in output_names ], - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, ) def to_uppercase(self: Self) -> ArrowExpr: + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: str(name).upper(), + alias_output_names=lambda output_names: [ + str(name).upper() for name in output_names + ], + ) + + def _from_colname_func_and_alias_output_names( + self: Self, + name_mapping_func: Callable[[str], str], + alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None, + ) -> ArrowExpr: return self._compliant_expr.__class__( - lambda df: [ - series.alias(str(name).upper()) + call=lambda df: [ + series.alias(name_mapping_func(name)) for series, name in zip( self._compliant_expr._call(df), self._compliant_expr._evaluate_output_names(df), @@ -123,9 +76,7 @@ def to_uppercase(self: Self) -> ArrowExpr: depth=self._compliant_expr._depth, function_name=self._compliant_expr._function_name, evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=lambda output_names: [ - str(name).upper() for name in output_names - ], + alias_output_names=alias_output_names, backend_version=self._compliant_expr._backend_version, version=self._compliant_expr._version, kwargs=self._compliant_expr._kwargs, diff --git a/narwhals/_dask/expr_name.py b/narwhals/_dask/expr_name.py index a261dcdf1..bf1ee32a6 100644 --- a/narwhals/_dask/expr_name.py +++ b/narwhals/_dask/expr_name.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING from typing import Callable +from typing import Sequence if TYPE_CHECKING: from typing_extensions import Self @@ -14,87 +15,53 @@ def __init__(self: Self, expr: DaskExpr) -> None: self._compliant_expr = expr def keep(self: Self) -> DaskExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=None, - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, - ) + return self._from_alias_output_names(alias_output_names=None) def map(self: Self, function: Callable[[str], str]) -> DaskExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ function(str(name)) for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "function": function}, ) def prefix(self: Self, prefix: str) -> DaskExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{prefix}{output_name}" for output_name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "prefix": prefix}, ) def suffix(self: Self, suffix: str) -> DaskExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{output_name}{suffix}" for output_name in output_names - ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "suffix": suffix}, + ] ) def to_lowercase(self: Self) -> DaskExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ str(name).lower() for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, ) def to_uppercase(self: Self) -> DaskExpr: + return self._from_alias_output_names( + alias_output_names=lambda output_names: [ + str(name).upper() for name in output_names + ] + ) + + def _from_alias_output_names( + self: Self, + alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None, + ) -> DaskExpr: return self._compliant_expr.__class__( - self._compliant_expr._call, + call=self._compliant_expr._call, depth=self._compliant_expr._depth, function_name=self._compliant_expr._function_name, evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=lambda output_names: [ - str(name).upper() for name in output_names - ], + alias_output_names=alias_output_names, returns_scalar=self._compliant_expr._returns_scalar, backend_version=self._compliant_expr._backend_version, version=self._compliant_expr._version, diff --git a/narwhals/_duckdb/expr_name.py b/narwhals/_duckdb/expr_name.py index a62933a20..e8f6a3142 100644 --- a/narwhals/_duckdb/expr_name.py +++ b/narwhals/_duckdb/expr_name.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING from typing import Callable +from typing import Sequence if TYPE_CHECKING: from typing_extensions import Self @@ -14,82 +15,53 @@ def __init__(self: Self, expr: DuckDBExpr) -> None: self._compliant_expr = expr def keep(self: Self) -> DuckDBExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=None, - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - ) + return self._from_alias_output_names(alias_output_names=None) def map(self: Self, function: Callable[[str], str]) -> DuckDBExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ function(str(name)) for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def prefix(self: Self, prefix: str) -> DuckDBExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{prefix}{output_name}" for output_name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def suffix(self: Self, suffix: str) -> DuckDBExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{output_name}{suffix}" for output_name in output_names - ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, + ] ) def to_lowercase(self: Self) -> DuckDBExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ str(name).lower() for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def to_uppercase(self: Self) -> DuckDBExpr: + return self._from_alias_output_names( + alias_output_names=lambda output_names: [ + str(name).upper() for name in output_names + ] + ) + + def _from_alias_output_names( + self: Self, + alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None, + ) -> DuckDBExpr: return self._compliant_expr.__class__( - self._compliant_expr._call, + call=self._compliant_expr._call, depth=self._compliant_expr._depth, function_name=self._compliant_expr._function_name, evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=lambda output_names: [ - str(name).upper() for name in output_names - ], + alias_output_names=alias_output_names, returns_scalar=self._compliant_expr._returns_scalar, backend_version=self._compliant_expr._backend_version, version=self._compliant_expr._version, diff --git a/narwhals/_pandas_like/expr_name.py b/narwhals/_pandas_like/expr_name.py index 065e4c34f..3c0c663a0 100644 --- a/narwhals/_pandas_like/expr_name.py +++ b/narwhals/_pandas_like/expr_name.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING from typing import Callable +from typing import Sequence if TYPE_CHECKING: from typing_extensions import Self @@ -14,112 +15,59 @@ def __init__(self: Self, expr: PandasLikeExpr) -> None: self._compliant_expr = expr def keep(self: Self) -> PandasLikeExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(name) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: name, alias_output_names=None, - implementation=self._compliant_expr._implementation, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, ) def map(self: Self, function: Callable[[str], str]) -> PandasLikeExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(function(str(name))) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: function(str(name)), alias_output_names=lambda output_names: [ function(str(name)) for name in output_names ], - implementation=self._compliant_expr._implementation, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "function": function}, ) def prefix(self: Self, prefix: str) -> PandasLikeExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(f"{prefix}{name}") - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: f"{prefix}{name}", alias_output_names=lambda output_names: [ f"{prefix}{output_name}" for output_name in output_names ], - implementation=self._compliant_expr._implementation, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "prefix": prefix}, ) def suffix(self: Self, suffix: str) -> PandasLikeExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(f"{name}{suffix}") - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: f"{name}{suffix}", alias_output_names=lambda output_names: [ f"{output_name}{suffix}" for output_name in output_names ], - implementation=self._compliant_expr._implementation, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs={**self._compliant_expr._kwargs, "suffix": suffix}, ) def to_lowercase(self: Self) -> PandasLikeExpr: - return self._compliant_expr.__class__( - lambda df: [ - series.alias(str(name).lower()) - for series, name in zip( - self._compliant_expr._call(df), - self._compliant_expr._evaluate_output_names(df), - ) - ], - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: str(name).lower(), alias_output_names=lambda output_names: [ str(name).lower() for name in output_names ], - implementation=self._compliant_expr._implementation, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - kwargs=self._compliant_expr._kwargs, ) def to_uppercase(self: Self) -> PandasLikeExpr: + return self._from_colname_func_and_alias_output_names( + name_mapping_func=lambda name: str(name).upper(), + alias_output_names=lambda output_names: [ + str(name).upper() for name in output_names + ], + ) + + def _from_colname_func_and_alias_output_names( + self: Self, + name_mapping_func: Callable[[str], str], + alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None, + ) -> PandasLikeExpr: return self._compliant_expr.__class__( - lambda df: [ - series.alias(str(name).upper()) + call=lambda df: [ + series.alias(name_mapping_func(name)) for series, name in zip( self._compliant_expr._call(df), self._compliant_expr._evaluate_output_names(df), @@ -128,11 +76,9 @@ def to_uppercase(self: Self) -> PandasLikeExpr: depth=self._compliant_expr._depth, function_name=self._compliant_expr._function_name, evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=lambda output_names: [ - str(name).upper() for name in output_names - ], - implementation=self._compliant_expr._implementation, + alias_output_names=alias_output_names, backend_version=self._compliant_expr._backend_version, + implementation=self._compliant_expr._implementation, version=self._compliant_expr._version, kwargs=self._compliant_expr._kwargs, ) diff --git a/narwhals/_spark_like/expr_name.py b/narwhals/_spark_like/expr_name.py index 76ce0e73d..3e2ef7dd1 100644 --- a/narwhals/_spark_like/expr_name.py +++ b/narwhals/_spark_like/expr_name.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING from typing import Callable +from typing import Sequence if TYPE_CHECKING: from typing_extensions import Self @@ -14,82 +15,53 @@ def __init__(self: Self, expr: SparkLikeExpr) -> None: self._compliant_expr = expr def keep(self: Self) -> SparkLikeExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=None, - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, - ) + return self._from_alias_output_names(alias_output_names=None) def map(self: Self, function: Callable[[str], str]) -> SparkLikeExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ function(str(name)) for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def prefix(self: Self, prefix: str) -> SparkLikeExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{prefix}{output_name}" for output_name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def suffix(self: Self, suffix: str) -> SparkLikeExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ f"{output_name}{suffix}" for output_name in output_names - ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, + ] ) def to_lowercase(self: Self) -> SparkLikeExpr: - return self._compliant_expr.__class__( - self._compliant_expr._call, - depth=self._compliant_expr._depth, - function_name=self._compliant_expr._function_name, - evaluate_output_names=self._compliant_expr._evaluate_output_names, + return self._from_alias_output_names( alias_output_names=lambda output_names: [ str(name).lower() for name in output_names ], - returns_scalar=self._compliant_expr._returns_scalar, - backend_version=self._compliant_expr._backend_version, - version=self._compliant_expr._version, ) def to_uppercase(self: Self) -> SparkLikeExpr: + return self._from_alias_output_names( + alias_output_names=lambda output_names: [ + str(name).upper() for name in output_names + ] + ) + + def _from_alias_output_names( + self: Self, + alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None, + ) -> SparkLikeExpr: return self._compliant_expr.__class__( self._compliant_expr._call, depth=self._compliant_expr._depth, function_name=self._compliant_expr._function_name, evaluate_output_names=self._compliant_expr._evaluate_output_names, - alias_output_names=lambda output_names: [ - str(name).upper() for name in output_names - ], + alias_output_names=alias_output_names, returns_scalar=self._compliant_expr._returns_scalar, backend_version=self._compliant_expr._backend_version, version=self._compliant_expr._version,