feat(python): Identify inefficient use of Python string `removeprefix…

…`, `removesuffix`, and `zfill` in `map_elements` (#19672)
pola-rs · Nov 7, 2024 · 8335f75 · 8335f75
1 parent d32cf0a
commit 8335f75
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git 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
@@ -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
     # ---------------------------------------------