From ae7343fc51233a7e9ceb0852a635129518cdefeb Mon Sep 17 00:00:00 2001 From: volaya Date: Thu, 4 Mar 2021 13:31:18 +0100 Subject: [PATCH 1/2] 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 --- pystac/extensions/base.py | 24 ++++++++++++++++++++ pystac/extensions/eo.py | 10 ++------- pystac/extensions/pointcloud.py | 30 +++++-------------------- pystac/extensions/projection.py | 40 +++++++-------------------------- pystac/extensions/timestamps.py | 11 +++------ pystac/extensions/view.py | 25 +++++---------------- 6 files changed, 48 insertions(+), 92 deletions(-) diff --git a/pystac/extensions/base.py b/pystac/extensions/base.py index 42d720f23..f11579fb4 100644 --- a/pystac/extensions/base.py +++ b/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): diff --git a/pystac/extensions/eo.py b/pystac/extensions/eo.py index 04d17c3d7..c0451a7c9 100644 --- a/pystac/extensions/eo.py +++ b/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 ''.format(self.item.id) diff --git a/pystac/extensions/pointcloud.py b/pystac/extensions/pointcloud.py index e0c467b51..f2ece2bdd 100644 --- a/pystac/extensions/pointcloud.py +++ b/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): diff --git a/pystac/extensions/projection.py b/pystac/extensions/projection.py index 0c8b8d604..a30ecec72 100644 --- a/pystac/extensions/projection.py +++ b/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): diff --git a/pystac/extensions/timestamps.py b/pystac/extensions/timestamps.py index 21f391057..d72ac746d 100644 --- a/pystac/extensions/timestamps.py +++ b/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): diff --git a/pystac/extensions/view.py b/pystac/extensions/view.py index bd0a2ae00..0b08244df 100644 --- a/pystac/extensions/view.py +++ b/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): From fbd300478b2c31767e1e02d7b64cbabd6a846a07 Mon Sep 17 00:00:00 2001 From: volaya Date: Thu, 4 Mar 2021 14:08:55 +0100 Subject: [PATCH 2/2] fixed test --- tests/extensions/test_timestamps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/extensions/test_timestamps.py b/tests/extensions/test_timestamps.py index 91084d0d6..7f484e477 100644 --- a/tests/extensions/test_timestamps.py +++ b/tests/extensions/test_timestamps.py @@ -50,7 +50,7 @@ def test_apply(self): self.assertIsNone(d) for p in ('expires', 'unpublished'): - self.assertIsNone(item.properties[p]) + self.assertNotIn(p, item.properties) def test_validate_timestamps(self): item = pystac.read_file(self.example_uri)