diff --git a/arpes/utilities/dict.py b/arpes/utilities/dict.py index 340f980b..729b160b 100644 --- a/arpes/utilities/dict.py +++ b/arpes/utilities/dict.py @@ -1,8 +1,9 @@ """Utilities for modifying, iterating over, and transforming dictionaries.""" + from __future__ import annotations import re -from typing import Any +from typing import Any, TypeVar from .xarray import lift_dataarray_attrs, lift_datavar_attrs @@ -45,14 +46,24 @@ def rename_keys( return d -def clean_keys(d: dict[str, Any]) -> dict[str, Any]: - """Renames dictionary keys so that they are more Pythonic.""" +T = TypeVar("T") + + +def clean_keys(d: dict[str, T]) -> dict[str, T]: + """Renames dict key to fit Pythonic more. + + Args: + d (dict): dictionary to be cleaned. + + Returns: + dict object whose key is cleaned. + """ def clean_single_key(k: str) -> str: k = k.replace(" ", "_") k = k.replace(".", "_") k = k.lower() - k = re.sub(r"[()/?]", "", k) + k = re.sub(r"[()/?]", "_", k) return k.replace("__", "_") return dict(zip([clean_single_key(k) for k in d], d.values(), strict=True)) diff --git a/tests/test_generic_utilities.py b/tests/test_generic_utilities.py index 505592b8..671d4413 100644 --- a/tests/test_generic_utilities.py +++ b/tests/test_generic_utilities.py @@ -1,7 +1,15 @@ """Test for generic utility.""" + import pytest -from arpes.utilities import deep_equals, deep_update +from arpes.utilities import clean_keys, deep_equals, deep_update + + +def test_cldean_keys() -> None: + """Test for clean_keys.""" + test_dict = {"Excitation Energy": 4.03, "Count/Cycle": 100} + cleaned_dict = clean_keys(test_dict) + assert cleaned_dict == {"excitation_energy": 4.03, "count_cycle": 100} @pytest.mark.parametrize( @@ -19,7 +27,7 @@ def test_deep_equals( *, expected_equal: bool, ) -> None: - """[TODO:summary]. + """Test for deep_equals. [TODO:description]