Extract method to correctly set value in Item/Asset for ItemExtensions

Refactored extension to use this parent method instead of locally implementing the same logic repeatedly

Author: volaya

pystac/extensions/base.py

@@ -121,6 +121,30 @@ def enable_extension(cls, stac_object):
         """
         pass
 
+    def _set_property(self, key, value, asset):
+        '''
+        Set an Item or an Asset property.
+
+        If an Asset is supplied, sets the property on the Asset.
+        Otherwise sets the Item's value.
+
+        If the passed value to set is None, the property key is removed from
+        the dictionary of properties.
+
+        It's recommended to use this method from extensions, instead of implementing
+        the logic for that in the corresponding subclasses.
+
+        Args:
+            key (str): The name of the property
+            value (Object): the value to set
+            asset: The Asset to modify. If None, the property will be set in the Item
+        '''
+        target = self.item.properties if asset is None else asset.properties
+        if value is None:
+            target.pop(key, None)
+        else:
+            target[key] = value
+
 
 class RegisteredSTACExtensions:
     def __init__(self, extension_definitions):

pystac/extensions/eo.py

@@ -85,10 +85,7 @@ def set_bands(self, bands, asset=None):
         Otherwise sets the Item's value.
         """
         band_dicts = [b.to_dict() for b in bands]
-        if asset is not None:
-            asset.properties['eo:bands'] = band_dicts
-        else:
-            self.item.properties['eo:bands'] = band_dicts
+        self._set_property('eo:bands', band_dicts, asset)
 
     @property
     def cloud_cover(self):
@@ -124,10 +121,7 @@ def set_cloud_cover(self, cloud_cover, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['eo:cloud_cover'] = cloud_cover
-        else:
-            asset.properties['eo:cloud_cover'] = cloud_cover
+        self._set_property('eo:cloud_cover', cloud_cover, asset)
 
     def __repr__(self):
         return '<EOItemExt Item id={}>'.format(self.item.id)

pystac/extensions/pointcloud.py

@@ -79,10 +79,7 @@ def set_count(self, count, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['pc:count'] = count
-        else:
-            asset.properties['pc:count'] = count
+        self._set_property('pc:count', count, asset)
 
     @property
     def type(self):
@@ -117,10 +114,7 @@ def set_type(self, type, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['pc:type'] = type
-        else:
-            asset.properties['pc:type'] = type
+        self._set_property('pc:type', type, asset)
 
     @property
     def encoding(self):
@@ -158,10 +152,7 @@ def set_encoding(self, encoding, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['pc:encoding'] = encoding
-        else:
-            asset.properties['pc:encoding'] = encoding
+        self._set_property('pc:encoding', encoding, asset)
 
     @property
     def schemas(self):
@@ -202,10 +193,7 @@ def set_schemas(self, schemas, asset=None):
         Otherwise sets the Item's value.
         """
         dicts = [s.to_dict() for s in schemas]
-        if asset is None:
-            self.item.properties['pc:schemas'] = dicts
-        else:
-            asset.properties['pc:schemas'] = dicts
+        self._set_property('pc:schemas', dicts, asset)
 
     @property
     def density(self):
@@ -242,10 +230,7 @@ def set_density(self, density, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['pc:density'] = density
-        else:
-            asset.properties['pc:density'] = density
+        self._set_property('pc:density', density, asset)
 
     @property
     def statistics(self):
@@ -289,10 +274,7 @@ def set_statistics(self, statistics, asset=None):
         """
         if statistics is not None:
             statistics = [s.to_dict() for s in statistics]
-        if asset is None:
-            self.item.properties['pc:statistics'] = statistics
-        else:
-            asset.properties['pc:statistics'] = statistics
+        self._set_property('pc:statistics', statistics, asset)
 
     @classmethod
     def _object_links(cls):

pystac/extensions/projection.py

@@ -103,10 +103,7 @@ def set_epsg(self, epsg, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:epsg'] = epsg
-        else:
-            asset.properties['proj:epsg'] = epsg
+        self._set_property('proj:epsg', epsg, asset)
 
     @property
     def wkt2(self):
@@ -147,10 +144,7 @@ def set_wkt2(self, wkt2, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:wkt2'] = wkt2
-        else:
-            asset.properties['proj:wkt2'] = wkt2
+        self._set_property('proj:wkt2', wkt2, asset)
 
     @property
     def projjson(self):
@@ -194,10 +188,7 @@ def set_projjson(self, projjson, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:projjson'] = projjson
-        else:
-            asset.properties['proj:projjson'] = projjson
+        self._set_property('proj:projjson', projjson, asset)
 
     @property
     def geometry(self):
@@ -239,10 +230,7 @@ def set_geometry(self, geometry, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:geometry'] = geometry
-        else:
-            asset.properties['proj:geometry'] = geometry
+        self._set_property('proj:geometry', geometry, asset)
 
     @property
     def bbox(self):
@@ -285,10 +273,7 @@ def set_bbox(self, bbox, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:bbox'] = bbox
-        else:
-            asset.properties['proj:bbox'] = bbox
+        self._set_property('proj:bbox', bbox, asset)
 
     @property
     def centroid(self):
@@ -330,10 +315,7 @@ def set_centroid(self, centroid, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:centroid'] = centroid
-        else:
-            asset.properties['proj:centroid'] = centroid
+        self._set_property('proj:centroid', centroid, asset)
 
     @property
     def shape(self):
@@ -373,10 +355,7 @@ def set_shape(self, shape, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:shape'] = shape
-        else:
-            asset.properties['proj:shape'] = shape
+        self._set_property('proj:shape', shape, asset)
 
     @property
     def transform(self):
@@ -419,10 +398,7 @@ def set_transform(self, transform, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['proj:transform'] = transform
-        else:
-            asset.properties['proj:transform'] = transform
+        self._set_property('proj:transform', transform, asset)
 
     @classmethod
     def _object_links(cls):

pystac/extensions/timestamps.py

@@ -66,14 +66,9 @@ def _timestamp_getter(self, key, asset=None):
         return timestamp
 
     def _timestamp_setter(self, timestamp, key, asset=None):
-        if timestamp is None:
-            self.item.properties[key] = timestamp
-        else:
-            timestamp_str = datetime_to_str(timestamp)
-            if asset is not None:
-                asset.properties[key] = timestamp_str
-            else:
-                self.item.properties[key] = timestamp_str
+        if timestamp is not None:
+            timestamp = datetime_to_str(timestamp)
+        self._set_property(key, timestamp, asset)
 
     @property
     def published(self):

pystac/extensions/view.py

@@ -101,10 +101,7 @@ def set_off_nadir(self, off_nadir, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['view:off_nadir'] = off_nadir
-        else:
-            asset.properties['view:off_nadir'] = off_nadir
+        self._set_property('view:off_nadir', off_nadir, asset)
 
     @property
     def incidence_angle(self):
@@ -141,10 +138,7 @@ def set_incidence_angle(self, incidence_angle, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['view:incidence_angle'] = incidence_angle
-        else:
-            asset.properties['view:incidence_angle'] = incidence_angle
+        self._set_property('view:incidence_angle', incidence_angle, asset)
 
     @property
     def azimuth(self):
@@ -181,10 +175,7 @@ def set_azimuth(self, azimuth, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['view:azimuth'] = azimuth
-        else:
-            asset.properties['view:azimuth'] = azimuth
+        self._set_property('view:azimuth', azimuth, asset)
 
     @property
     def sun_azimuth(self):
@@ -220,10 +211,7 @@ def set_sun_azimuth(self, sun_azimuth, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['view:sun_azimuth'] = sun_azimuth
-        else:
-            asset.properties['view:sun_azimuth'] = sun_azimuth
+        self._set_property('view:sun_azimuth', sun_azimuth, asset)
 
     @property
     def sun_elevation(self):
@@ -259,10 +247,7 @@ def set_sun_elevation(self, sun_elevation, asset=None):
         If an Asset is supplied, sets the property on the Asset.
         Otherwise sets the Item's value.
         """
-        if asset is None:
-            self.item.properties['view:sun_elevation'] = sun_elevation
-        else:
-            asset.properties['view:sun_elevation'] = sun_elevation
+        self._set_property('view:sun_elevation', sun_elevation, asset)
 
     @classmethod
     def _object_links(cls):