diff --git a/src/uiprotect/__init__.py b/src/uiprotect/__init__.py index 4172de01..85f3d41e 100644 --- a/src/uiprotect/__init__.py +++ b/src/uiprotect/__init__.py @@ -7,7 +7,6 @@ from .utils import ( get_nested_attr, get_nested_attr_as_bool, - get_top_level_attr, get_top_level_attr_as_bool, make_enabled_getter, make_required_getter, @@ -21,7 +20,6 @@ "ProtectApiClient", "get_nested_attr", "get_nested_attr_as_bool", - "get_top_level_attr", "get_top_level_attr_as_bool", "make_value_getter", "make_enabled_getter", diff --git a/src/uiprotect/utils.py b/src/uiprotect/utils.py index 90aa20a4..f8a99ae8 100644 --- a/src/uiprotect/utils.py +++ b/src/uiprotect/utils.py @@ -634,7 +634,7 @@ def get_nested_attr(attrs: tuple[str, ...], obj: Any) -> Any: for key in attrs: if (value := getattr(value, key, _SENTINEL)) is _SENTINEL: return None - return value.value if isinstance(value, Enum) else value + return value def get_nested_attr_as_bool(attrs: tuple[str, ...], obj: Any) -> bool: @@ -643,25 +643,18 @@ def get_nested_attr_as_bool(attrs: tuple[str, ...], obj: Any) -> bool: for key in attrs: if (value := getattr(value, key, _SENTINEL)) is _SENTINEL: return False - return bool(value.value if isinstance(value, Enum) else value) - - -def get_top_level_attr(attr: str, obj: Any) -> Any: - """Fetch a top level attribute.""" - value = getattr(obj, attr) - return value.value if isinstance(value, Enum) else value + return bool(value) def get_top_level_attr_as_bool(attr: str, obj: Any) -> Any: """Fetch a top level attribute as a bool.""" - value = getattr(obj, attr) - return bool(value.value if isinstance(value, Enum) else value) + return bool(getattr(obj, attr)) def make_value_getter(ufp_value: str) -> Callable[[T], Any]: """Return a function to get a value from a Protect device.""" if "." not in ufp_value: - return partial(get_top_level_attr, ufp_value) + return attrgetter(ufp_value) return partial(get_nested_attr, tuple(ufp_value.split("."))) diff --git a/tests/test_utils.py b/tests/test_utils.py index 72c83f28..e1ddfc74 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -16,7 +16,6 @@ dict_diff, get_nested_attr, get_nested_attr_as_bool, - get_top_level_attr, get_top_level_attr_as_bool, make_enabled_getter, make_required_getter, @@ -219,7 +218,7 @@ def test_get_nested_attr(): data = Mock(a=Mock(b=Mock(c=1)), d=3, f=_MockEnum.C) assert get_nested_attr(("a", "b", "c"), data) == 1 assert get_nested_attr(("d",), data) == 3 - assert get_nested_attr(("f",), data) == _MockEnum.C.value + assert get_nested_attr(("f",), data) == _MockEnum.C @pytest.mark.asyncio @@ -230,18 +229,9 @@ def test_get_nested_attr_as_bool(): assert get_nested_attr_as_bool(("f",), data) is True -@pytest.mark.asyncio -def test_get_top_level_attr(): - data = Mock(a=1, b=2, c=3, d=_MockEnum.C) - assert get_top_level_attr("a", data) == 1 - assert get_top_level_attr("b", data) == 2 - assert get_top_level_attr("c", data) == 3 - assert get_top_level_attr("d", data) == _MockEnum.C.value - - @pytest.mark.asyncio def test_get_top_level_attr_as_bool(): - data = Mock(a=True, b=False, c=_MockEnum.C, d=None) + data = Mock(a=True, b=False, c=True, d=None) assert get_top_level_attr_as_bool("a", data) is True assert get_top_level_attr_as_bool("b", data) is False assert get_top_level_attr_as_bool("c", data) is True @@ -250,11 +240,11 @@ def test_get_top_level_attr_as_bool(): @pytest.mark.asyncio def test_make_value_getter(): - data = Mock(a=1, b=2, c=3, d=_MockEnum.C) + data = Mock(a=1, b=2, c=3, d=4) assert make_value_getter("a")(data) == 1 assert make_value_getter("b")(data) == 2 assert make_value_getter("c")(data) == 3 - assert make_value_getter("d")(data) == _MockEnum.C.value + assert make_value_getter("d")(data) == 4 @pytest.mark.asyncio