Use enums #194

fastkml/kml.py

@@ -40,7 +40,11 @@
 from fastkml.base import _BaseObject
 from fastkml.data import ExtendedData
 from fastkml.data import Schema
+from fastkml.enums import GridOrigin
+from fastkml.enums import RefreshMode
+from fastkml.enums import Shape
 from fastkml.enums import Verbosity
+from fastkml.enums import ViewRefreshMode
 from fastkml.geometry import AnyGeometryType
 from fastkml.geometry import LinearRing
 from fastkml.geometry import LineString
@@ -529,9 +533,9 @@ class Icon(_BaseObject):
     __name__ = "Icon"
 
     _href = None
-    _refresh_mode: str = None
+    _refresh_mode: Optional[RefreshMode]
     _refresh_interval = None
-    _view_refresh_mode = None
+    _view_refresh_mode: Optional[ViewRefreshMode]
     _view_refresh_time = None
     _view_bound_scale = None
     _view_format = None
@@ -543,9 +547,9 @@ def __init__(
         id: Optional[str] = None,
         target_id: Optional[str] = None,
         href: Optional[str] = None,
-        refresh_mode: Optional[str] = None,
+        refresh_mode: Optional[RefreshMode] = None,
         refresh_interval: Optional[float] = None,
-        view_refresh_mode: Optional[str] = None,
+        view_refresh_mode: Optional[ViewRefreshMode] = None,
         view_refresh_time: Optional[float] = None,
         view_bound_scale: Optional[float] = None,
         view_format: Optional[str] = None,
@@ -577,7 +581,7 @@ def href(self, href) -> None:
             raise ValueError
 
     @property
-    def refresh_mode(self) -> Optional[str]:
+    def refresh_mode(self) -> Optional[RefreshMode]:
         """
         Specifies a time-based refresh mode.
 
@@ -590,13 +594,8 @@ def refresh_mode(self) -> Optional[str]:
         return self._refresh_mode
 
     @refresh_mode.setter
-    def refresh_mode(self, refresh_mode) -> None:
-        if isinstance(refresh_mode, str):
-            self._refresh_mode = refresh_mode
-        elif refresh_mode is None:
-            self._refresh_mode = None
-        else:
-            raise ValueError
+    def refresh_mode(self, refresh_mode: Optional[RefreshMode]) -> None:
+        self._refresh_mode = refresh_mode
 
     @property
     def refresh_interval(self) -> Optional[float]:
@@ -613,7 +612,7 @@ def refresh_interval(self, refresh_interval: Optional[float]) -> None:
             raise ValueError
 
     @property
-    def view_refresh_mode(self):
+    def view_refresh_mode(self) -> Optional[ViewRefreshMode]:
         """
         Specifies how the link is refreshed when the "camera" changes.
 
@@ -625,18 +624,13 @@ def view_refresh_mode(self):
          - onRequest - Refresh the file only when the user explicitly requests it.
            (For example, in Google Earth, the user right-clicks and selects Refresh
             in the Context menu.)
-          - onRegion - Refresh the file when the Region becomes active.
+         - onRegion - Refresh the file when the Region becomes active.
         """
         return self._view_refresh_mode
 
     @view_refresh_mode.setter
-    def view_refresh_mode(self, view_refresh_mode) -> None:
-        if isinstance(view_refresh_mode, str):
-            self._view_refresh_mode = view_refresh_mode
-        elif view_refresh_mode is None:
-            self._view_refresh_mode = None
-        else:
-            raise ValueError
+    def view_refresh_mode(self, view_refresh_mode: Optional[ViewRefreshMode]) -> None:
+        self._view_refresh_mode = view_refresh_mode
 
     @property
     def view_refresh_time(self):
@@ -749,7 +743,7 @@ def etree_element(
             href.text = self._href
         if self._refresh_mode:
             refresh_mode = config.etree.SubElement(element, f"{self.ns}refreshMode")
-            refresh_mode.text = self._refresh_mode
+            refresh_mode.text = self._refresh_mode.value
         if self._refresh_interval:
             refresh_interval = config.etree.SubElement(
                 element,
@@ -761,7 +755,7 @@ def etree_element(
                 element,
                 f"{self.ns}viewRefreshMode",
             )
-            view_refresh_mode.text = self._view_refresh_mode
+            view_refresh_mode.text = self._view_refresh_mode.value
         if self._view_refresh_time:
             view_refresh_time = config.etree.SubElement(
                 element,
@@ -792,7 +786,7 @@ def from_element(self, element: Element) -> None:
 
         refresh_mode = element.find(f"{self.ns}refreshMode")
         if refresh_mode is not None:
-            self.refresh_mode = refresh_mode.text
+            self.refresh_mode = RefreshMode(refresh_mode.text)
 
         refresh_interval = element.find(f"{self.ns}refreshInterval")
         if refresh_interval is not None:
@@ -803,7 +797,7 @@ def from_element(self, element: Element) -> None:
 
         view_refresh_mode = element.find(f"{self.ns}viewRefreshMode")
         if view_refresh_mode is not None:
-            self.view_refresh_mode = view_refresh_mode.text
+            self.view_refresh_mode = ViewRefreshMode(view_refresh_mode.text)
 
         view_refresh_time = element.find(f"{self.ns}viewRefreshTime")
         if view_refresh_time is not None:
@@ -1089,7 +1083,7 @@ class PhotoOverlay(_Overlay):
     _max_height = None
     # Height in pixels of the original image.
 
-    _grid_origin = None
+    _grid_origin: Optional[GridOrigin]
     # Specifies where to begin numbering the tiles in each layer of the pyramid.
     # A value of lowerLeft specifies that row 1, column 1 of each layer is in
     # the bottom left corner of the grid.
@@ -1100,7 +1094,7 @@ class PhotoOverlay(_Overlay):
     # The icon drawn is specified by the <styleUrl> and <StyleSelector> fields,
     # just as it is for <Placemark>.
 
-    _shape = "rectangle"
+    _shape: Optional[Shape]
     # The PhotoOverlay is projected onto the <shape>.
     # The <shape> can be one of the following:
     #   rectangle (default) -
@@ -1228,17 +1222,12 @@ def max_height(self, value) -> None:
             raise ValueError
 
     @property
-    def grid_origin(self) -> Optional[str]:
+    def grid_origin(self) -> Optional[GridOrigin]:
         return self._grid_origin
 
     @grid_origin.setter
-    def grid_origin(self, value) -> None:
-        if value in ("lowerLeft", "upperLeft"):
-            self._grid_origin = str(value)
-        elif value is None:
-            self._grid_origin = None
-        else:
-            raise ValueError
+    def grid_origin(self, value: Optional[GridOrigin]) -> None:
+        self._grid_origin = value
 
     @property
     def point(self) -> str:
@@ -1252,18 +1241,12 @@ def point(self, value) -> None:
             raise ValueError
 
     @property
-    def shape(self) -> Optional[str]:
+    def shape(self) -> Optional[Shape]:
         return self._shape
 
     @shape.setter
-    def shape(self, value) -> None:
-        if value in ("rectangle", "cylinder", "sphere"):
-            self._shape = str(value)
-        elif value is None:
-            self._shape = None
-        else:
-            msg = "Shape must be one of rectangle, cylinder, sphere"
-            raise ValueError(msg)
+    def shape(self, value: Optional[Shape]) -> None:
+        self._shape = value
 
     def view_volume(self, left_fov, right_fov, bottom_fov, top_fov, near) -> None:
         self.left_fov = left_fov

fastkml/styles.py

@@ -34,6 +34,7 @@
 from fastkml import config
 from fastkml.base import _BaseObject
 from fastkml.enums import ColorMode
+from fastkml.enums import DisplayMode
 from fastkml.enums import Units
 from fastkml.enums import Verbosity
 from fastkml.types import Element
@@ -131,7 +132,7 @@ class _ColorStyle(_BaseObject):
     # The order of expression is aabbggrr, where aa=alpha (00 to ff);
     # bb=blue (00 to ff); gg=green (00 to ff); rr=red (00 to ff).
 
-    color_mode: ColorMode
+    color_mode: Optional[ColorMode]
     # Values for <colorMode> are normal (no effect) and random.
     # A value of random applies a random linear scale to the base <color>
 
@@ -307,10 +308,10 @@ def _get_kwargs(
                 kwargs["icon_href"] = href.text
         hot_spot = element.find(f"{ns}hotSpot")
         if hot_spot is not None:
-            x = hot_spot.attrib.get("x")
-            y = hot_spot.attrib.get("y")
-            xunits = hot_spot.attrib.get("xunits")
-            yunits = hot_spot.attrib.get("yunits")
+            x = hot_spot.attrib.get("x")  # type: ignore[attr-defined]
+            y = hot_spot.attrib.get("y")  # type: ignore[attr-defined]
+            xunits = hot_spot.attrib.get("xunits")  # type: ignore[attr-defined]
+            yunits = hot_spot.attrib.get("yunits")  # type: ignore[attr-defined]
             x = float(x) if x is not None else 0
             y = float(y) if y is not None else 0
             xunits = Units(xunits) if xunits is not None else None
@@ -412,7 +413,7 @@ def __init__(
         id: Optional[str] = None,
         target_id: Optional[str] = None,
         color: Optional[str] = None,
-        color_mode: Optional[str] = None,
+        color_mode: Optional[ColorMode] = None,
         fill: int = 1,
         outline: int = 1,
     ) -> None:
@@ -579,7 +580,7 @@ class BalloonStyle(_BaseObject):
     # in the Feature elements that use this BalloonStyle:
     # <text>This is $[name], whose description is:<br/>$[description]</text>
 
-    display_mode = None
+    display_mode: Optional[DisplayMode]
     # If <displayMode> is default, Google Earth uses the information supplied
     # in <text> to create a balloon . If <displayMode> is hide, Google Earth
     # does not display the balloon. In Google Earth, clicking the List View
@@ -595,7 +596,7 @@ def __init__(
         bg_color: Optional[str] = None,
         text_color: Optional[str] = None,
         text: Optional[str] = None,
-        display_mode: Optional[str] = None,
+        display_mode: Optional[DisplayMode] = None,
     ) -> None:
         super().__init__(ns=ns, name_spaces=name_spaces, id=id, target_id=target_id)
         self.bg_color = bg_color
@@ -632,7 +633,7 @@ def etree_element(
                 element,
                 f"{self.ns}displayMode",
             )
-            elem.text = self.display_mode
+            elem.text = self.display_mode.value
         return element
 
     @classmethod
@@ -665,7 +666,7 @@ def _get_kwargs(
             kwargs["text"] = text.text
         display_mode = element.find(f"{ns}displayMode")
         if display_mode is not None:
-            kwargs["display_mode"] = display_mode.text
+            kwargs["display_mode"] = DisplayMode(display_mode.text)
         return kwargs
 
 

tests/kml_test.py

@@ -16,6 +16,8 @@
 
 """Test the kml classes."""
 from fastkml import kml
+from fastkml.enums import RefreshMode
+from fastkml.enums import ViewRefreshMode
 from tests.base import Lxml
 from tests.base import StdLibrary
 
@@ -67,9 +69,9 @@ def test_icon_read(self) -> None:
 
         assert icon.id == "icon-01"
         assert icon.href == "http://maps.google.com/mapfiles/kml/paddle/red-circle.png"
-        assert icon.refresh_mode == "onInterval"
+        assert icon.refresh_mode == RefreshMode("onInterval")
         assert icon.refresh_interval == 60
-        assert icon.view_refresh_mode == "onStop"
+        assert icon.view_refresh_mode == ViewRefreshMode("onStop")
         assert icon.view_refresh_time == 4
         assert icon.view_bound_scale == 1.2
         assert icon.view_format == "BBOX=[bboxWest],[bboxSouth],[bboxEast],[bboxNorth]"

tests/oldunit_test.py

@@ -27,6 +27,7 @@
 from fastkml import styles
 from fastkml.enums import AltitudeMode
 from fastkml.enums import ColorMode
+from fastkml.enums import DisplayMode
 
 try:
     import lxml
@@ -812,7 +813,7 @@ def test_balloonstyle(self) -> None:
         assert isinstance(style, styles.BalloonStyle)
         assert style.bg_color == "ffffffbb"
         assert style.text_color == "ff000000"
-        assert style.display_mode == "default"
+        assert style.display_mode == DisplayMode.default
         assert "$[geDirections]" in style.text
         assert "$[description]" in style.text
         k2 = kml.KML()

tests/styles_test.py

@@ -18,6 +18,7 @@
 
 from fastkml import styles
 from fastkml.enums import ColorMode
+from fastkml.enums import DisplayMode
 from tests.base import Lxml
 from tests.base import StdLibrary
 
@@ -206,7 +207,7 @@ def test_balloon_style(self) -> None:
             bg_color="7fff0000",
             text_color="ff00ff00",
             text="<b>Hello</b>",
-            display_mode="hide",
+            display_mode=DisplayMode.hide,
         )
 
         serialized = bs.to_string()
@@ -238,7 +239,7 @@ def test_balloon_style_read(self) -> None:
         assert bs.bg_color == "7fff1144"
         assert bs.text_color == "ff11ff22"
         assert bs.text == "<b>World</b>"
-        assert bs.display_mode == "default"
+        assert bs.display_mode == DisplayMode.default
 
     def test_style(self) -> None:
         icons = styles.IconStyle(
@@ -263,7 +264,7 @@ def test_style(self) -> None:
             bg_color="7fff0000",
             text_color="ff00ff00",
             text="<b>Hello</b>",
-            display_mode="hide",
+            display_mode=DisplayMode.hide,
         )
         ls = styles.LabelStyle(
             id="id-a0",
@@ -347,7 +348,7 @@ def test_style_read(self) -> None:
         assert next(iter(style.styles())).bg_color == "7fff0000"  # type: ignore[union-attr]
         assert next(iter(style.styles())).text_color == "ff00ff00"  # type: ignore[union-attr]
         assert next(iter(style.styles())).text == "<b>Hello</b>"  # type: ignore[union-attr]
-        assert next(iter(style.styles())).display_mode == "hide"  # type: ignore[union-attr]
+        assert next(iter(style.styles())).display_mode == DisplayMode.hide  # type: ignore[union-attr]
 
         assert list(style.styles())[1].id == "id-i0"
         assert list(style.styles())[1].target_id == "target-i0"
@@ -400,7 +401,7 @@ def test_stylemap(self) -> None:
             bg_color="7fff0000",
             text_color="ff00ff00",
             text="<b>Hello</b>",
-            display_mode="hide",
+            display_mode=DisplayMode.hide,
         )
         ls = styles.LabelStyle(
             id="id-a0",