Skip to content

Commit

Permalink
force LabelEnums to have labels, gives .label type str causing less m…
Browse files Browse the repository at this point in the history
…ypy headache

- update LabelEnum to require label and make description default to empty str for better usability
- update tests to cover new LabelEnum behavior
  • Loading branch information
janosh committed Feb 8, 2025
1 parent 067c63d commit 9bebc60
Show file tree
Hide file tree
Showing 5 changed files with 12 additions and 25 deletions.
8 changes: 4 additions & 4 deletions pymatviz/coordination/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,10 @@
class CnSplitMode(LabelEnum):
"""How to split the coordination number histogram into subplots."""

none = "none"
by_element = "by element"
by_structure = "by structure"
by_structure_and_element = "by structure and element"
none = "none", "None"
by_element = "by element", "By element"
by_structure = "by structure", "By structure"
by_structure_and_element = "by structure and element", "By structure and element"


def create_hover_text(
Expand Down
15 changes: 5 additions & 10 deletions pymatviz/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,12 +85,7 @@ class LabelEnum(StrEnum):
Simply add label and description as a tuple starting with the key's value.
"""

def __new__(
cls,
val: str,
label: str | None = None,
desc: str | None = None,
) -> Self:
def __new__(cls, val: str, label: str, desc: str = "") -> Self:
"""Create a new class from a value, label, and description. Label and
description are optional.
"""
Expand All @@ -111,14 +106,14 @@ def __reduce_ex__(self, proto: object) -> tuple[type, tuple[str]]:
return str, (self.value,)

@property
def label(self) -> str | None:
def label(self) -> str:
"""Make label read-only."""
return self.__dict__.get("label")
return self.__dict__["label"]

@property
def description(self) -> str | None:
def description(self) -> str:
"""Make description read-only."""
return self.__dict__.get("desc")
return self.__dict__["desc"]


eV_per_atom = html_tag("(eV/atom)", style="small") # noqa: N816
Expand Down
7 changes: 0 additions & 7 deletions pymatviz/process_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from __future__ import annotations

import itertools
import warnings
from typing import TYPE_CHECKING

import pandas as pd
Expand Down Expand Up @@ -62,12 +61,6 @@ def count_elements(
Returns:
pd.Series: Map element symbols to heatmap values.
"""
# TODO: remove warning after 2025-01-01
if fill_value is None:
warnings.warn(
"default value of fill_value changed from zero to None.", stacklevel=2
)

valid_count_modes = set(ElemCountMode)
if count_mode not in valid_count_modes:
raise ValueError(f"Invalid {count_mode=} must be one of {valid_count_modes}")
Expand Down
3 changes: 2 additions & 1 deletion pymatviz/scatter.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,8 @@ def density_scatter_plotly(
y (str): y-values dataframe column name.
df (pd.DataFrame): DataFrame with x and y columns.
density ('kde' | 'interpolate' | 'empirical'): Determines the method for
calculating and displaying density.
calculating and displaying density. Default is 'empirical' when n_bins
is provided, else 'kde' for kernel density estimation.
log_density (bool | None): Whether to apply logarithmic scaling to density.
If None, automatically set based on density range.
identity_line (bool | dict[str, Any], optional): Whether to add a parity line
Expand Down
4 changes: 1 addition & 3 deletions tests/test_enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

class DummyEnum(LabelEnum):
TEST1 = "test1", "Test Label 1", "Test Description 1"
TEST2 = "test2"
TEST2 = "test2", "Test Label 2"


def test_str_enum() -> None:
Expand Down Expand Up @@ -74,8 +74,6 @@ def test_label_enum_new() -> None:

def test_label_enum_repr() -> None:
assert repr(DummyEnum.TEST1) == DummyEnum.TEST1.label == "Test Label 1"
assert repr(DummyEnum.TEST2) == "DummyEnum.TEST2"
assert DummyEnum.TEST2.label is None


@pytest.mark.parametrize(
Expand Down

0 comments on commit 9bebc60

Please sign in to comment.