aboutsummaryrefslogtreecommitdiff
path: root/venv/lib/python3.8/site-packages/narwhals/expr_name.py
blob: 4f77bec0a53f205cfdec60b1e7af9f0c05dbed79 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
from __future__ import annotations

from typing import TYPE_CHECKING, Callable, Generic, TypeVar

if TYPE_CHECKING:
    from narwhals.expr import Expr

ExprT = TypeVar("ExprT", bound="Expr")


class ExprNameNamespace(Generic[ExprT]):
    def __init__(self, expr: ExprT) -> None:
        self._expr = expr

    def keep(self) -> ExprT:
        r"""Keep the original root name of the expression.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import pandas as pd
            >>> import narwhals as nw
            >>> df_native = pd.DataFrame({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> df.select(nw.col("foo").alias("alias_for_foo").name.keep()).columns
            ['foo']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.keep()
        )

    def map(self, function: Callable[[str], str]) -> ExprT:
        r"""Rename the output of an expression by mapping a function over the root name.

        Arguments:
            function: Function that maps a root name to a new name.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import pandas as pd
            >>> import narwhals as nw
            >>> df_native = pd.DataFrame({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> renaming_func = lambda s: s[::-1]  # reverse column name
            >>> df.select(nw.col("foo", "BAR").name.map(renaming_func)).columns
            ['oof', 'RAB']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.map(function)
        )

    def prefix(self, prefix: str) -> ExprT:
        r"""Add a prefix to the root column name of the expression.

        Arguments:
            prefix: Prefix to add to the root column name.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import polars as pl
            >>> import narwhals as nw
            >>> df_native = pl.DataFrame({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> df.select(nw.col("foo", "BAR").name.prefix("with_prefix")).columns
            ['with_prefixfoo', 'with_prefixBAR']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.prefix(prefix)
        )

    def suffix(self, suffix: str) -> ExprT:
        r"""Add a suffix to the root column name of the expression.

        Arguments:
            suffix: Suffix to add to the root column name.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import polars as pl
            >>> import narwhals as nw
            >>> df_native = pl.DataFrame({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> df.select(nw.col("foo", "BAR").name.suffix("_with_suffix")).columns
            ['foo_with_suffix', 'BAR_with_suffix']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.suffix(suffix)
        )

    def to_lowercase(self) -> ExprT:
        r"""Make the root column name lowercase.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import pyarrow as pa
            >>> import narwhals as nw
            >>> df_native = pa.table({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> df.select(nw.col("foo", "BAR").name.to_lowercase()).columns
            ['foo', 'bar']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.to_lowercase()
        )

    def to_uppercase(self) -> ExprT:
        r"""Make the root column name uppercase.

        Returns:
            A new expression.

        Notes:
            This will undo any previous renaming operations on the expression.
            Due to implementation constraints, this method can only be called as the last
            expression in a chain. Only one name operation per expression will work.

        Examples:
            >>> import pyarrow as pa
            >>> import narwhals as nw
            >>> df_native = pa.table({"foo": [1, 2], "BAR": [4, 5]})
            >>> df = nw.from_native(df_native)
            >>> df.select(nw.col("foo", "BAR").name.to_uppercase()).columns
            ['FOO', 'BAR']
        """
        return self._expr._with_elementwise_op(
            lambda plx: self._expr._to_compliant_expr(plx).name.to_uppercase()
        )