Skip to content

Commit 5539471

Browse files
committed
Use enum for recommended task values
1 parent 1415f95 commit 5539471

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

pystac/extensions/label.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,18 @@ def __str__(self) -> str:
5252
"""Convenience attribute for checking if values are valid label types"""
5353

5454

55+
class LabelTask(str, Enum):
56+
"""Enumerates recommended label tasks."""
57+
58+
def __str__(self) -> str:
59+
return str(self.value)
60+
61+
REGRESSION = "regression"
62+
CLASSIFICATION = "classification"
63+
DETECTION = "detection"
64+
SEGMENTATION = "segmentation"
65+
66+
5567
class LabelClasses:
5668
"""Defines the list of possible class names (e.g., tree, building, car, hippo).
5769
@@ -403,7 +415,7 @@ def apply(
403415
label_type: LabelType,
404416
label_properties: Optional[List[str]] = None,
405417
label_classes: Optional[List[LabelClasses]] = None,
406-
label_tasks: Optional[List[str]] = None,
418+
label_tasks: Optional[List[Union[LabelTask, str]]] = None,
407419
label_methods: Optional[List[str]] = None,
408420
label_overviews: Optional[List[LabelOverview]] = None,
409421
) -> None:
@@ -499,14 +511,14 @@ def label_classes(self, v: Optional[List[LabelClasses]]) -> None:
499511
self.obj.properties[CLASSES_PROP] = classes
500512

501513
@property
502-
def label_tasks(self) -> Optional[List[str]]:
514+
def label_tasks(self) -> Optional[List[Union[LabelTask, str]]]:
503515
"""Gets or set a list of tasks these labels apply to. Usually a subset of 'regression',
504516
'classification', 'detection', or 'segmentation', but may be arbitrary
505517
values."""
506518
return self.obj.properties.get(TASKS_PROP)
507519

508520
@label_tasks.setter
509-
def label_tasks(self, v: Optional[List[str]]) -> None:
521+
def label_tasks(self, v: Optional[List[Union[LabelTask, str]]]) -> None:
510522
if v is None:
511523
self.obj.properties.pop(TASKS_PROP, None)
512524
else:

0 commit comments

Comments
 (0)