Merge pull request #261 from kylebarron/kyle/use-enum

Subclass relevant classes from Enum

Author: lossyrob

CHANGELOG.md

@@ -6,6 +6,10 @@
 
 - Fix handling of optional properties when using apply on view extension ([#259](https://github.com/stac-utils/pystac/pull/259))
 
+### Changed
+
+- Subclass relevant classes from `enum.Enum`. This allows iterating over the class' contents. The `__str__` method is overwritten so this should not break backwards compatibility. ([#261](https://github.com/stac-utils/pystac/pull/261))
+
 ## [v0.5.4]
 
 ### Added

pystac/catalog.py

@@ -1,5 +1,6 @@
 import os
 from copy import deepcopy
+from enum import Enum
 
 import pystac
 from pystac import STACError
@@ -10,7 +11,10 @@
 from pystac.utils import (is_absolute_href, make_absolute_href)
 
 
-class CatalogType:
+class CatalogType(str, Enum):
+    def __str__(self):
+        return str(self.value)
+
     SELF_CONTAINED = 'SELF_CONTAINED'
     """A 'self-contained catalog' is one that is designed for portability.
     Users may want to download a catalog from online and be able to use it on their
@@ -38,8 +42,8 @@ class CatalogType:
         `The best practices documentation on published catalogs <https://github.com/radiantearth/stac-spec/blob/v0.8.1/best-practices.md#published-catalogs>`_
     """ # noqa E501
 
-    @staticmethod
-    def determine_type(stac_json):
+    @classmethod
+    def determine_type(cls, stac_json):
         """Determines the catalog type based on a STAC JSON dict.
 
         Only applies to Catalogs or Collections
@@ -61,12 +65,12 @@ def determine_type(stac_json):
 
         if self_link:
             if relative:
-                return CatalogType.RELATIVE_PUBLISHED
+                return cls.RELATIVE_PUBLISHED
             else:
-                return CatalogType.ABSOLUTE_PUBLISHED
+                return cls.ABSOLUTE_PUBLISHED
         else:
             if relative:
-                return CatalogType.SELF_CONTAINED
+                return cls.SELF_CONTAINED
             else:
                 return None
 

pystac/extensions/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa
+from enum import Enum
 
 
 class ExtensionError(Exception):
@@ -7,8 +8,11 @@ class ExtensionError(Exception):
     pass
 
 
-class Extensions:
+class Extensions(str, Enum):
     """Enumerates the IDs of common extensions."""
+    def __str__(self):
+        return str(self.value)
+
     CHECKSUM = 'checksum'
     COLLECTION_ASSETS = 'collection-assets'
     DATACUBE = 'datacube'

pystac/extensions/label.py

@@ -1,14 +1,19 @@
 """STAC Model classes for Label extension.
 """
+from enum import Enum
+
 from pystac import STACError
 from pystac.extensions import Extensions
 from pystac.extensions.base import (ItemExtension, ExtensionDefinition, ExtendedObject)
 from pystac.item import (Item, Asset)
 from pystac.link import Link
 
 
-class LabelType:
+class LabelType(str, Enum):
     """Enumerates valid label types (RASTER or VECTOR)."""
+    def __str__(self):
+        return str(self.value)
+
     VECTOR = 'vector'
     RASTER = 'raster'
 

pystac/link.py

@@ -1,12 +1,16 @@
 from copy import (copy, deepcopy)
+from enum import Enum
 
 from pystac import STACError
 from pystac.stac_io import STAC_IO
 from pystac.utils import (make_absolute_href, make_relative_href, is_absolute_href)
 
 
-class LinkType:
+class LinkType(str, Enum):
     """Enumerates link types; used to determine if a link is absolute or relative."""
+    def __str__(self):
+        return str(self.value)
+
     ABSOLUTE = 'ABSOLUTE'
     RELATIVE = 'RELATIVE'
 

pystac/media_type.py

@@ -1,6 +1,12 @@
-class MediaType:
+from enum import Enum
+
+
+class MediaType(str, Enum):
     """A list of common media types that can be used in STAC Asset and Link metadata.
     """
+    def __str__(self):
+        return str(self.value)
+
     COG = 'image/tiff; application=geotiff; profile=cloud-optimized'
     GEOJSON = 'application/geo+json'
     GEOPACKAGE = 'application/geopackage+sqlite3'

pystac/stac_object.py

@@ -1,4 +1,5 @@
 from abc import (ABC, abstractmethod)
+from enum import Enum
 
 import pystac
 from pystac import STACError
@@ -8,7 +9,10 @@
 from pystac.extensions import ExtensionError
 
 
-class STACObjectType:
+class STACObjectType(str, Enum):
+    def __str__(self):
+        return str(self.value)
+
     CATALOG = 'CATALOG'
     COLLECTION = 'COLLECTION'
     ITEM = 'ITEM'