1
1
from dataclasses import dataclass
2
2
from typing import Any , Generic , Literal , TypeVar , cast
3
3
4
- from pystac import Asset , Catalog , Collection , Item , STACError
4
+ from pystac import Asset , Catalog , Collection , Item , Link , STACError
5
5
from pystac .extensions .classification import ClassificationExtension
6
6
from pystac .extensions .datacube import DatacubeExtension
7
7
from pystac .extensions .eo import EOExtension
22
22
from pystac .extensions .view import ViewExtension
23
23
from pystac .extensions .xarray_assets import XarrayAssetsExtension
24
24
25
- T = TypeVar ("T" , Asset , AssetDefinition )
25
+ T = TypeVar ("T" , Asset , AssetDefinition , Link )
26
+ U = TypeVar ("U" , Asset , AssetDefinition )
26
27
27
28
EXTENSION_NAMES = Literal [
28
29
"classification" ,
@@ -196,14 +197,14 @@ def xarray(self) -> XarrayAssetsExtension[Item]:
196
197
return XarrayAssetsExtension .ext (self .stac_object )
197
198
198
199
199
- class _AssetExt (Generic [T ]):
200
+ class _AssetsExt (Generic [T ]):
200
201
stac_object : T
201
202
202
203
def has (self , name : EXTENSION_NAMES ) -> bool :
203
204
if self .stac_object .owner is None :
204
205
raise STACError (
205
- f"Attempted to add extension='{ name } ' for an Asset with no owner. "
206
- "Use Asset .set_owner and then try to add the extension again."
206
+ f"Attempted to add extension='{ name } ' for an object with no owner. "
207
+ "Use ` .set_owner` and then try to add the extension again."
207
208
)
208
209
else :
209
210
return cast (
@@ -213,67 +214,71 @@ def has(self, name: EXTENSION_NAMES) -> bool:
213
214
def add (self , name : EXTENSION_NAMES ) -> None :
214
215
if self .stac_object .owner is None :
215
216
raise STACError (
216
- f"Attempted to add extension='{ name } ' for an Asset with no owner. "
217
- "Use Asset .set_owner and then try to add the extension again."
217
+ f"Attempted to add extension='{ name } ' for an object with no owner. "
218
+ "Use ` .set_owner` and then try to add the extension again."
218
219
)
219
220
else :
220
221
_get_class_by_name (name ).add_to (self .stac_object .owner )
221
222
222
223
def remove (self , name : EXTENSION_NAMES ) -> None :
223
224
if self .stac_object .owner is None :
224
225
raise STACError (
225
- f"Attempted to remove extension='{ name } ' for an Asset with no owner. "
226
- "Use Asset .set_owner and then try to remove the extension again."
226
+ f"Attempted to remove extension='{ name } ' for an object with no owner. "
227
+ "Use ` .set_owner` and then try to remove the extension again."
227
228
)
228
229
else :
229
230
_get_class_by_name (name ).remove_from (self .stac_object .owner )
230
231
232
+
233
+ class _AssetExt (_AssetsExt [U ]):
234
+ stac_object : U
235
+
231
236
@property
232
- def classification (self ) -> ClassificationExtension [T ]:
237
+ def classification (self ) -> ClassificationExtension [U ]:
233
238
return ClassificationExtension .ext (self .stac_object )
234
239
235
240
@property
236
- def cube (self ) -> DatacubeExtension [T ]:
241
+ def cube (self ) -> DatacubeExtension [U ]:
237
242
return DatacubeExtension .ext (self .stac_object )
238
243
239
244
@property
240
- def eo (self ) -> EOExtension [T ]:
245
+ def eo (self ) -> EOExtension [U ]:
241
246
return EOExtension .ext (self .stac_object )
242
247
243
248
@property
244
- def pc (self ) -> PointcloudExtension [T ]:
249
+ def pc (self ) -> PointcloudExtension [U ]:
245
250
return PointcloudExtension .ext (self .stac_object )
246
251
247
252
@property
248
- def proj (self ) -> ProjectionExtension [T ]:
253
+ def proj (self ) -> ProjectionExtension [U ]:
249
254
return ProjectionExtension .ext (self .stac_object )
250
255
251
256
@property
252
- def raster (self ) -> RasterExtension [T ]:
257
+ def raster (self ) -> RasterExtension [U ]:
253
258
return RasterExtension .ext (self .stac_object )
254
259
255
260
@property
256
- def sar (self ) -> SarExtension [T ]:
261
+ def sar (self ) -> SarExtension [U ]:
257
262
return SarExtension .ext (self .stac_object )
258
263
259
264
@property
260
- def sat (self ) -> SatExtension [T ]:
265
+ def sat (self ) -> SatExtension [U ]:
261
266
return SatExtension .ext (self .stac_object )
262
267
263
268
@property
264
- def storage (self ) -> StorageExtension [T ]:
269
+ def storage (self ) -> StorageExtension [U ]:
265
270
return StorageExtension .ext (self .stac_object )
266
271
267
272
@property
268
- def table (self ) -> TableExtension [T ]:
273
+ def table (self ) -> TableExtension [U ]:
269
274
return TableExtension .ext (self .stac_object )
270
275
271
276
@property
272
- def version (self ) -> BaseVersionExtension [T ]:
277
+ def version (self ) -> BaseVersionExtension [U ]:
273
278
return BaseVersionExtension .ext (self .stac_object )
274
279
275
280
@property
276
- def view (self ) -> ViewExtension [T ]:
281
+ def view (self ) -> ViewExtension [U ]:
277
282
return ViewExtension .ext (self .stac_object )
278
283
279
284
@@ -282,7 +287,7 @@ class AssetExt(_AssetExt[Asset]):
282
287
stac_object : Asset
283
288
284
289
@property
285
- def file (self ) -> FileExtension :
290
+ def file (self ) -> FileExtension [ Asset ] :
286
291
return FileExtension .ext (self .stac_object )
287
292
288
293
@property
@@ -297,3 +302,12 @@ def xarray(self) -> XarrayAssetsExtension[Asset]:
297
302
@dataclass
298
303
class ItemAssetExt (_AssetExt [AssetDefinition ]):
299
304
stac_object : AssetDefinition
305
+
306
+
307
+ @dataclass
308
+ class LinkExt (_AssetsExt [Link ]):
309
+ stac_object : Link
310
+
311
+ @property
312
+ def file (self ) -> FileExtension [Link ]:
313
+ return FileExtension .ext (self .stac_object )
0 commit comments