Add enum_attribute

Author: srittau

NEWS.md

@@ -1,14 +1,10 @@
-# News in version 2.1.0
-
-## API Additions
-
-* Add `autocomplete` attribute to `Input`, `TextArea`, `Select`, and `Form`.
-
 # News in version 2.0.0
 
 ## API Additions
 
 * Add `Button.disabled`.
+* Add `autocomplete` attribute to `Input`, `TextArea`, `Select`, and `Form`.
+* Add `enum_attribute`.
 
 ## Incompatible Changes
 

htmlgen/attribute.py

@@ -1,4 +1,5 @@
 import datetime
+from enum import Enum
 
 from htmlgen.timeutil import parse_rfc3339_partial_time
 
@@ -273,3 +274,55 @@ def __set__(self, obj, value):
             obj.add_css_classes(self._css_class)
         else:
             obj.remove_css_classes(self._css_class)
+
+
+class enum_attribute:
+    """Add an attribute to an HTML element that only allows limited values.
+
+        >>> from enum import Enum
+        >>> from htmlgen import Element
+        >>> class MyEnum(Enum):
+        ...     FOO = "foo"
+        ...     BAR = "bar"
+        >>> class MyElement(Element):
+        ...     value = enum_attribute("value", MyEnum)
+        >>> element = MyElement("div")
+        >>> element.value
+        >>> str(element)
+        '<div></div>'
+        >>> element.value = MyEnum.FOO
+        >>> str(element)
+        '<div value="foo"></div>'
+
+    If the optional default argument is given, the attribute will not be
+    included if the value matches it.
+
+        >>> class MyElement(Element):
+        ...     value = enum_attribute("value", MyEnum, default=MyEnum.FOO)
+        >>> element = MyElement("div")
+        >>> element.value
+        MyEnum.FOO
+        >>> str(element)
+        '<div></div>'
+    """
+
+    def __init__(self, attribute_name, enum, default=None):
+        if not issubclass(enum, Enum):
+            raise TypeError("enum must be an Enum class")
+        self._attribute_name = attribute_name
+        self._enum = enum
+        self._default = default
+
+    def __get__(self, obj, _=None):
+        value = obj.get_attribute(self._attribute_name, None)
+        if value is None:
+            return self._default
+        return self._enum(value)
+
+    def __set__(self, obj, value):
+        if value is None:
+            obj.remove_attribute(self._attribute_name)
+        elif not isinstance(value, self._enum):
+            raise TypeError("value must be an {}".format(self._enum))
+        else:
+            obj.set_attribute(self._attribute_name, value.value)

htmlgen/attribute.pyi

@@ -1,8 +1,11 @@
 import datetime
-from typing import Optional, List, Iterable
+from enum import Enum
+from typing import Generic, Iterable, List, Optional, Type, TypeVar
 
 from htmlgen.element import Element
 
+_E = TypeVar("_E", bound=Enum)
+
 class html_attribute:
     def __init__(
         self, attribute_name: str, default: Optional[str] = ...
@@ -62,3 +65,12 @@ class css_class_attribute:
     def __init__(self, css_class: str) -> None: ...
     def __get__(self, obj: Element, type_: Optional[type] = ...) -> bool: ...
     def __set__(self, obj: Element, value: bool) -> None: ...
+
+class enum_attribute(Generic[_E]):
+    def __init__(
+        self, attribute_name: str, enum: Type[_E], default: Optional[_E] = ...
+    ) -> None: ...
+    def __get__(
+        self, obj: Element, type: Optional[type] = ...
+    ) -> Optional[_E]: ...
+    def __set__(self, obj: Element, value: Optional[_E]) -> None: ...

test_htmlgen/attribute.py

@@ -1,7 +1,9 @@
 import datetime
+from enum import Enum
 from unittest import TestCase
 
-from asserts import assert_true, assert_false, assert_is_none, assert_equal
+from asserts import assert_true, assert_false, assert_is_none, assert_equal, \
+    assert_raises
 
 from htmlgen.attribute import (
     html_attribute,
@@ -12,6 +14,7 @@
     list_html_attribute,
     data_attribute,
     css_class_attribute,
+    enum_attribute,
 )
 from htmlgen.element import Element
 
@@ -226,3 +229,51 @@ class MyElement(Element):
         assert_true(element.has_css_class("my-class"))
         element.attr = True
         assert_true(element.has_css_class("my-class"))
+
+
+class TestEnum(Enum):
+    FOO = "foo"
+    BAR = "bar"
+
+
+class EnumAttributeTest(TestCase):
+    def test_enum(self):
+        class MyElement(Element):
+            attr = enum_attribute("attr", TestEnum)
+
+        element = MyElement("div")
+        assert_is_none(element.attr)
+        assert_equal("<div></div>", str(element))
+        element.attr = TestEnum.BAR
+        assert_equal(TestEnum.BAR, element.attr)
+        assert_equal('<div attr="bar"></div>', str(element))
+        element.attr = None
+        assert_is_none(element.attr)
+        assert_equal('<div></div>', str(element))
+
+    def test_default(self):
+        class MyElement(Element):
+            attr = enum_attribute("attr", TestEnum, default=TestEnum.FOO)
+
+        element = MyElement("div")
+        assert_equal(TestEnum.FOO, element.attr)
+        assert_equal("<div></div>", str(element))
+        element.attr = TestEnum.BAR
+        assert_equal(TestEnum.BAR, element.attr)
+        assert_equal('<div attr="bar"></div>', str(element))
+        element.attr = None
+        assert_equal(TestEnum.FOO, element.attr)
+        assert_equal("<div></div>", str(element))
+
+    def test_not_an_enum(self):
+        with assert_raises(TypeError):
+            class MyElement(Element):
+                attr = enum_attribute("attr", "foo")  # type: ignore
+
+    def test_invalid_value(self):
+        class MyElement(Element):
+            attr = enum_attribute("attr", TestEnum)
+
+        element = MyElement("div")
+        with assert_raises(TypeError):
+            element.attr = "foo"  # type: ignore