From ccaf682eaa42725831ce461ad31158327e184b8a Mon Sep 17 00:00:00 2001 From: Alexander Beedie Date: Fri, 29 Nov 2024 12:05:48 +0400 Subject: [PATCH] feat(python): Improved error message on invalid Python `Enum` init (#20060) --- py-polars/polars/convert/general.py | 2 +- py-polars/polars/datatypes/classes.py | 6 +++--- py-polars/tests/unit/datatypes/test_enum.py | 4 +--- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/py-polars/polars/convert/general.py b/py-polars/polars/convert/general.py index 1776adf1cf6a..13256d925c18 100644 --- a/py-polars/polars/convert/general.py +++ b/py-polars/polars/convert/general.py @@ -782,7 +782,7 @@ def _from_series_repr(m: re.Match[str]) -> Series: else: string_values = [ v.strip() - for v in re.findall(r"[\s>#]*(?:\t|\s{4,})([^\n]*)\n", m.groups()[-1]) + for v in re.findall(r"[\s>#]*(?:\t|\s{2,})([^\n]*)\n", m.groups()[-1]) ] if string_values == ["[", "]"]: string_values = [] diff --git a/py-polars/polars/datatypes/classes.py b/py-polars/polars/datatypes/classes.py index 4008775072ac..4e022a12205e 100644 --- a/py-polars/polars/datatypes/classes.py +++ b/py-polars/polars/datatypes/classes.py @@ -622,10 +622,10 @@ def __init__(self, categories: Series | Iterable[str] | type[enum.Enum]) -> None ) if isclass(categories) and issubclass(categories, enum.Enum): - for enum_subclass in (enum.IntFlag, enum.Flag, enum.IntEnum): + for enum_subclass in (enum.Flag, enum.IntEnum): if issubclass(categories, enum_subclass): - enum_type_name = enum_subclass.__name__ - msg = f"Enum categories must be strings; Python `enum.{enum_type_name}` values are integers" + enum_type_name = categories.__name__ + msg = f"Enum categories must be strings; `{enum_type_name}` values are integers" raise TypeError(msg) enum_values = [ diff --git a/py-polars/tests/unit/datatypes/test_enum.py b/py-polars/tests/unit/datatypes/test_enum.py index bd23ddcf623e..9830faaf7380 100644 --- a/py-polars/tests/unit/datatypes/test_enum.py +++ b/py-polars/tests/unit/datatypes/test_enum.py @@ -94,11 +94,9 @@ class Color(EnumBase): # type: ignore[no-redef,misc,valid-type] GREEN = enum.auto() BLUE = enum.auto() - base_name = EnumBase.__name__ - with pytest.raises( TypeError, - match=f"Enum categories must be strings; Python `enum.{base_name}` values are integers", + match="Enum categories must be strings; `Color` values are integers", ): pl.Enum(Color)