From 8335f75c80bf1dc04b8cbb6ec91052ea747f45fa Mon Sep 17 00:00:00 2001 From: Alexander Beedie Date: Thu, 7 Nov 2024 20:16:34 +0400 Subject: [PATCH] feat(python): Identify inefficient use of Python string `removeprefix`, `removesuffix`, and `zfill` in `map_elements` (#19672) --- py-polars/polars/_utils/udfs.py | 3 +++ .../operations/map/test_inefficient_map_warning.py | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/py-polars/polars/_utils/udfs.py b/py-polars/polars/_utils/udfs.py index cfc9d06492b1..ed91c1920cc9 100644 --- a/py-polars/polars/_utils/udfs.py +++ b/py-polars/polars/_utils/udfs.py @@ -183,12 +183,15 @@ class OpNames: "endswith": "str.ends_with", "lower": "str.to_lowercase", "lstrip": "str.strip_chars_start", + "removeprefix": "str.strip_prefix", + "removesuffix": "str.strip_suffix", "replace": "str.replace", "rstrip": "str.strip_chars_end", "startswith": "str.starts_with", "strip": "str.strip_chars", "title": "str.to_titlecase", "upper": "str.to_uppercase", + "zfill": "str.zfill", # temporal "date": "dt.date", "isoweekday": "dt.weekday", diff --git a/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py b/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py index 13eee2c731e3..61416054755d 100644 --- a/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py +++ b/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py @@ -186,6 +186,16 @@ "lambda x: x.replace(':','',2)", """pl.col("b").str.replace(':','',n=2,literal=True)""", ), + ( + "b", + "lambda x: x.removeprefix('A').removesuffix('F')", + """pl.col("b").str.strip_prefix('A').str.strip_suffix('F')""", + ), + ( + "b", + "lambda x: x.zfill(8)", + """pl.col("b").str.zfill(8)""", + ), # --------------------------------------------- # json expr: load/extract # ---------------------------------------------