Use an enum for predicate types

Signed-off-by: Facundo Tuesca <[email protected]>
sigstore · Sep 13, 2024 · 2b0584b · 2b0584b
1 parent 65f16bf
commit 2b0584b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/sigstore/_cli.py b/sigstore/_cli.py
@@ -42,9 +42,7 @@
 from sigstore._utils import sha256_digest
 from sigstore.dsse import StatementBuilder, Subject
 from sigstore.dsse._predicate import (
-    SUPPORTED_PREDICATE_TYPES,
-    PREDICATE_TYPE_SLSA_v0_2,
-    PREDICATE_TYPE_SLSA_v1_0,
+    PredicateType,
     SLSAPredicateV0_2,
     SLSAPredicateV1_0,
 )
@@ -277,10 +275,10 @@ def _parser() -> argparse.ArgumentParser:
     dsse_options.add_argument(
         "--predicate-type",
         metavar="TYPE",
-        choices=SUPPORTED_PREDICATE_TYPES,
+        choices=[p.value for p in PredicateType],
         type=str,
         required=True,
-        help=f"Specify a predicate type ({', '.join(SUPPORTED_PREDICATE_TYPES)})",
+        help=f"Specify a predicate type ({', '.join(list(PredicateType))})",
     )
 
     oidc_options = attest.add_argument_group("OpenID Connect options")
@@ -704,17 +702,17 @@ def _attest(args: argparse.Namespace) -> None:
             # Since most of the predicate fields are optional, this only checks that
             # the fields that are present and correctly spelled have the expected
             # type.
-            if args.predicate_type == PREDICATE_TYPE_SLSA_v0_2:
+            if args.predicate_type == PredicateType.SLSA_v0_2:
                 SLSAPredicateV0_2.model_validate(predicate)
-            elif args.predicate_type == PREDICATE_TYPE_SLSA_v1_0:
+            elif args.predicate_type == PredicateType.SLSA_v1_0:
                 SLSAPredicateV1_0.model_validate(predicate)
             else:
                 _invalid_arguments(
                     args,
-                    f'Unsupported predicate type "{args.predicate_type}". Predicate type must be one of: {SUPPORTED_PREDICATE_TYPES}',
+                    f'Unsupported predicate type "{args.predicate_type}". Predicate type must be one of: {list(PredicateType)}',
                 )
 
-    except ValidationError as e:
+    except (ValidationError, json.JSONDecodeError) as e:
         _invalid_arguments(
             args, f'Unable to parse predicate of type "{args.predicate_type}": {e}'
         )

diff --git a/sigstore/dsse/_predicate.py b/sigstore/dsse/_predicate.py
@@ -16,6 +16,7 @@
 Models for the predicates used in in-toto statements
 """
 
+import enum
 from typing import Any, Dict, List, Literal, Optional, TypeVar, Union
 
 from pydantic import (
@@ -30,10 +31,15 @@
 
 from sigstore.dsse import Digest
 
-PREDICATE_TYPE_SLSA_v0_2 = "https://slsa.dev/provenance/v0.2"
-PREDICATE_TYPE_SLSA_v1_0 = "https://slsa.dev/provenance/v1"
 
-SUPPORTED_PREDICATE_TYPES = [PREDICATE_TYPE_SLSA_v0_2, PREDICATE_TYPE_SLSA_v1_0]
+class PredicateType(str, enum.Enum):
+    """
+    Currently supported predicate types
+    """
+
+    SLSA_v0_2 = "https://slsa.dev/provenance/v0.2"
+    SLSA_v1_0 = "https://slsa.dev/provenance/v1"
+
 
 # Common models
 SourceDigest = Union[Literal["sha1"], Literal["gitCommit"]]