EdgeType and ArgRole enums

telmomenezes · telmomenezes · commit e04c5bfc0b5c · 2026-04-07T15:01:29.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### Added
 - [] pattern notation for specifying sequences of arguments.
+- EdgeType and ArgRole enums.
 
 ### Changed
 - multiple patterns functions are now Hyperedge/Atom methods: is_wildcard, is_pattern, is_fun_pattern, is_variable, contains_variable, variable_name
diff --git a/src/hyperbase/cli/repl.py b/src/hyperbase/cli/repl.py
@@ -20,6 +20,7 @@
 from rich.tree import Tree
 
 from hyperbase.builders import hedge
+from hyperbase.constants import EdgeType
 from hyperbase.hyperedge import Atom, Hyperedge
 from hyperbase.parsers import Parser, get_parser, list_parsers
 from hyperbase.parsers.correctness import badness_check
@@ -42,14 +43,14 @@
 
 
 TYPE_COLORS = {
-    "C": "#4A9EFF",
-    "M": "#00E5CC",
-    "B": "#5DC4FF",
-    "P": "#FF8C42",
-    "T": "#00CDB8",
-    "J": "#00FF87",
-    "R": "#FFD700",
-    "S": "#FF6EC7",
+    EdgeType.CONCEPT: "#4A9EFF",
+    EdgeType.MODIFIER: "#00E5CC",
+    EdgeType.BUILDER: "#5DC4FF",
+    EdgeType.PREDICATE: "#FF8C42",
+    EdgeType.TRIGGER: "#00CDB8",
+    EdgeType.CONJUNCTION: "#00FF87",
+    EdgeType.RELATION: "#FFD700",
+    EdgeType.SPECIFIER: "#FF6EC7",
 }
 
 
diff --git a/src/hyperbase/constants.py b/src/hyperbase/constants.py
@@ -1,3 +1,36 @@
+from enum import Enum
+
+
+class EdgeType(str, Enum):
+    """SH edge types."""
+
+    CONCEPT = "C"
+    PREDICATE = "P"
+    MODIFIER = "M"
+    BUILDER = "B"
+    TRIGGER = "T"
+    CONJUNCTION = "J"
+    RELATION = "R"
+    SPECIFIER = "S"
+
+
+class ArgRole(str, Enum):
+    """Argument roles for predicates and builders."""
+
+    MAIN = "m"
+    SUBJECT = "s"
+    PASSIVE = "p"
+    AGENT = "a"
+    COMPLEMENT = "c"
+    OBJECT = "o"
+    INDIRECT = "i"
+    PARATAXIS = "t"
+    INTERJECTION = "j"
+    SPECIFICATION = "x"
+    RELATIVE = "r"
+    UNDETERMINED = "?"
+
+
 # Pre-defined system atoms
 compound_noun_builder = "+/B/."
 possessive_builder = "poss/Bp.am/."
@@ -8,20 +41,32 @@
 
 # Argument role ordering for normalisation
 argrole_order: dict[str, int] = {
-    "m": -1,
-    "s": 0,
-    "p": 1,
-    "a": 2,
-    "c": 3,
-    "o": 4,
-    "i": 5,
-    "t": 6,
-    "j": 7,
-    "x": 8,
-    "r": 9,
-    "?": 10,
+    ArgRole.MAIN: -1,
+    ArgRole.SUBJECT: 0,
+    ArgRole.PASSIVE: 1,
+    ArgRole.AGENT: 2,
+    ArgRole.COMPLEMENT: 3,
+    ArgRole.OBJECT: 4,
+    ArgRole.INDIRECT: 5,
+    ArgRole.PARATAXIS: 6,
+    ArgRole.INTERJECTION: 7,
+    ArgRole.SPECIFICATION: 8,
+    ArgRole.RELATIVE: 9,
+    ArgRole.UNDETERMINED: 10,
 }
 
 # Valid argument roles by connector type
-valid_p_argroles: set[str] = {"s", "p", "a", "c", "o", "i", "t", "j", "x", "r", "?"}
-valid_b_argroles: set[str] = {"m", "a"}
+valid_p_argroles: set[str] = {
+    ArgRole.SUBJECT,
+    ArgRole.PASSIVE,
+    ArgRole.AGENT,
+    ArgRole.COMPLEMENT,
+    ArgRole.OBJECT,
+    ArgRole.INDIRECT,
+    ArgRole.PARATAXIS,
+    ArgRole.INTERJECTION,
+    ArgRole.SPECIFICATION,
+    ArgRole.RELATIVE,
+    ArgRole.UNDETERMINED,
+}
+valid_b_argroles: set[str] = {ArgRole.MAIN, ArgRole.AGENT}
diff --git a/src/hyperbase/correctness.py b/src/hyperbase/correctness.py
@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import hyperbase.constants as const
+from hyperbase.constants import EdgeType
 
 if TYPE_CHECKING:
     from hyperbase.hyperedge import Atom, Hyperedge
@@ -21,7 +22,14 @@ def _check_atom(atom: Atom) -> dict[Hyperedge, list[tuple[str, str]]]:
     errors: list[tuple[str, str]] = []
 
     at = atom.mtype()
-    if at not in {"C", "P", "M", "B", "T", "J"}:
+    if at not in {
+        EdgeType.CONCEPT,
+        EdgeType.PREDICATE,
+        EdgeType.MODIFIER,
+        EdgeType.BUILDER,
+        EdgeType.TRIGGER,
+        EdgeType.CONJUNCTION,
+    }:
         errors.append(("bad-atom-type", f"{at} is not a valid atom type"))
 
     if len(errors) > 0:
@@ -36,49 +44,55 @@ def _check_edge(edge: Hyperedge) -> dict[Hyperedge, list[tuple[str, str]]]:
 
     ct = edge[0].mtype()
     # check if connector has valid type
-    if ct not in {"P", "M", "B", "T", "J"}:
+    if ct not in {
+        EdgeType.PREDICATE,
+        EdgeType.MODIFIER,
+        EdgeType.BUILDER,
+        EdgeType.TRIGGER,
+        EdgeType.CONJUNCTION,
+    }:
         errors.append(("conn-bad-type", f"connector has incorrect type: {ct}"))
     # check if modifier structure is correct
-    if ct == "M":
+    if ct == EdgeType.MODIFIER:
         if len(edge) != 2:
             errors.append(("mod-1-arg", "modifiers can only have one argument"))
     # check if builder structure is correct
-    elif ct == "B":
+    elif ct == EdgeType.BUILDER:
         if len(edge) != 3:
             errors.append(("build-2-args", "builders can only have two arguments"))
         for arg in edge[1:]:
             at = arg.mtype()
-            if at != "C":
+            if at != EdgeType.CONCEPT:
                 e = f"builder argument {arg!s} has incorrect type: {at}"
                 errors.append(("build-arg-bad-type", e))
     # check if trigger structure is correct
-    elif ct == "T":
+    elif ct == EdgeType.TRIGGER:
         if len(edge) != 2:
             errors.append(("trig-1-arg", "triggers can only have one arguments"))
         for arg in edge[1:]:
             at = arg.mtype()
-            if at not in {"C", "R"}:
+            if at not in {EdgeType.CONCEPT, EdgeType.RELATION}:
                 e = f"trigger argument {arg!s} has incorrect type: {at}"
                 errors.append(("trig-bad-arg-type", e))
     # check if predicate structure is correct
-    elif ct == "P":
+    elif ct == EdgeType.PREDICATE:
         for arg in edge[1:]:
             at = arg.mtype()
-            if at not in {"C", "R", "S"}:
+            if at not in {EdgeType.CONCEPT, EdgeType.RELATION, EdgeType.SPECIFIER}:
                 e = f"predicate argument {arg!s} has incorrect type: {at}"
                 errors.append(("pred-arg-bad-type", e))
     # check if conjunction structure is correct
-    elif ct == "J" and len(edge) < 3:
+    elif ct == EdgeType.CONJUNCTION and len(edge) < 3:
         errors.append(
             ("conj-2-args-min", "conjunctions must have at least two arguments")
         )
 
     # check argrole counts
-    if ct in {"P", "B"}:
+    if ct in {EdgeType.PREDICATE, EdgeType.BUILDER}:
         try:
             ars = edge.argroles()
             if len(ars) > 0:
-                if ct == "P":
+                if ct == EdgeType.PREDICATE:
                     for ar in ars:
                         if ar not in const.valid_p_argroles:
                             errors.append(
@@ -88,7 +102,7 @@ def _check_edge(edge: Hyperedge) -> dict[Hyperedge, list[tuple[str, str]]]:
                                     "for connector of type P",
                                 )
                             )
-                elif ct == "B":
+                elif ct == EdgeType.BUILDER:
                     for ar in ars:
                         if ar not in const.valid_b_argroles:
                             errors.append(
diff --git a/src/hyperbase/hyperedge.py b/src/hyperbase/hyperedge.py
@@ -4,6 +4,8 @@
 from dataclasses import dataclass
 from typing import Any, overload
 
+from hyperbase.constants import EdgeType
+
 
 @dataclass(frozen=True, init=False, eq=False, repr=False)
 class Hyperedge:
@@ -222,19 +224,19 @@ def type(self) -> str:
         Type inference is performed.
         """
         ptype = self[0].type()
-        if ptype[0] == "P":
-            outter_type = "R"
-        elif ptype[0] == "M":
+        if ptype[0] == EdgeType.PREDICATE:
+            outter_type = EdgeType.RELATION
+        elif ptype[0] == EdgeType.MODIFIER:
             if len(self) < 2:
                 raise RuntimeError(
                     f"Edge is malformed, type cannot be determined: {self!s}"
                 )
             return self[1].type()  # type: ignore[no-any-return]
-        elif ptype[0] == "T":
-            outter_type = "S"
-        elif ptype[0] == "B":
-            outter_type = "C"
-        elif ptype[0] == "J":
+        elif ptype[0] == EdgeType.TRIGGER:
+            outter_type = EdgeType.SPECIFIER
+        elif ptype[0] == EdgeType.BUILDER:
+            outter_type = EdgeType.CONCEPT
+        elif ptype[0] == EdgeType.CONJUNCTION:
             if len(self) < 2:
                 raise RuntimeError(
                     f"Edge is malformed, type cannot be determined: {self!s}"
@@ -245,7 +247,7 @@ def type(self) -> str:
                 f"Edge is malformed, type cannot be determined: {self!s}"
             )
 
-        return f"{outter_type}{ptype[1:]}"
+        return outter_type + ptype[1:]
 
     def connector_type(self) -> str | None:
         """Returns the type of the edge's connector.
@@ -306,9 +308,12 @@ def argroles(self) -> str:
         of/B.ma has argument roles "ma".
         """
         et = self.mtype()
-        if et in {"R", "C"} and self[0].mtype() in {"B", "P"}:
+        if et in {EdgeType.RELATION, EdgeType.CONCEPT} and self[0].mtype() in {
+            EdgeType.BUILDER,
+            EdgeType.PREDICATE,
+        }:
             return self[0].argroles()  # type: ignore[no-any-return]
-        if et not in {"B", "P"}:
+        if et not in {EdgeType.BUILDER, EdgeType.PREDICATE}:
             return ""
         return self[1].argroles()  # type: ignore[no-any-return]
 
@@ -552,7 +557,7 @@ def atom_with_type(self, atom_type: str) -> Atom | None:
 
     def argroles(self) -> str:
         et = self.mtype()
-        if et not in {"B", "P"}:
+        if et not in {EdgeType.BUILDER, EdgeType.PREDICATE}:
             return ""
         role = self.role()
         if len(role) < 2:
diff --git a/src/hyperbase/parsers/correctness.py b/src/hyperbase/parsers/correctness.py
@@ -2,6 +2,7 @@
 
 from typing import Any
 
+from hyperbase.constants import EdgeType
 from hyperbase.hyperedge import Hyperedge
 from hyperbase.parsers.utils import filter_alphanumeric_strings
 
@@ -36,9 +37,13 @@ def _visit(current_edge: Hyperedge) -> None:
 
         # Junction checks
         try:
-            if current_edge[0].mt == "J":
+            if current_edge[0].mt == EdgeType.CONJUNCTION:
                 types = {child.mt for child in current_edge[1:]}
-                if types != {"R"} and types != {"C"} and types != {"R", "S"}:
+                if (
+                    types != {EdgeType.RELATION}
+                    and types != {EdgeType.CONCEPT}
+                    and types != {EdgeType.RELATION, EdgeType.SPECIFIER}
+                ):
                     current_errors.append(
                         (
                             "bad-junction-types",
diff --git a/src/hyperbase/patterns/matcher.py b/src/hyperbase/patterns/matcher.py
@@ -6,6 +6,7 @@
 
 import hyperbase.constants as const
 from hyperbase import hedge
+from hyperbase.constants import EdgeType
 from hyperbase.hyperedge import Atom, Hyperedge
 
 # tok_pos can be nested lists/ints matching the edge structure
@@ -180,7 +181,7 @@ def _matches_atomic_pattern(edge: Hyperedge, atomic_pattern: Atom) -> bool:
                 return False
 
             # argroles match
-            if ap_type[0] in {"B", "P"}:
+            if ap_type[0] in {EdgeType.BUILDER, EdgeType.PREDICATE}:
                 ap_argroles_parts = ap_role[1].split("-")
                 if len(ap_argroles_parts) == 1:
                     ap_argroles_parts.append("")
diff --git a/src/hyperbase/transforms.py b/src/hyperbase/transforms.py
@@ -4,14 +4,15 @@
 
 import hyperbase.constants as const
 from hyperbase.builders import hedge
+from hyperbase.constants import EdgeType
 from hyperbase.hyperedge import Atom, Hyperedge, UniqueAtom
 
 
 def normalise(edge: Hyperedge) -> Hyperedge:
     """Return a normalised copy of the edge (argument roles sorted)."""
     if edge.atom:
         atom = cast(Atom, edge)
-        if atom.mtype() in {"B", "P"}:
+        if atom.mtype() in {EdgeType.BUILDER, EdgeType.PREDICATE}:
             ar = atom.argroles()
             if len(ar) > 0:
                 if ar[0] == "{":
@@ -104,11 +105,11 @@ def replace_argroles(edge: Hyperedge, argroles: str | None) -> Hyperedge:
         return Atom("/".join(parts))
     else:
         st = edge.mtype()
-        if st in {"C", "R"}:
+        if st in {EdgeType.CONCEPT, EdgeType.RELATION}:
             new_edge = [replace_argroles(edge[0], argroles)]
             new_edge += edge[1:]
             return Hyperedge(new_edge)
-        elif st in {"P", "B"}:
+        elif st in {EdgeType.PREDICATE, EdgeType.BUILDER}:
             new_edge = [edge[0], replace_argroles(edge[1], argroles)]
             new_edge += list(edge[2:])
             return Hyperedge(new_edge)
@@ -132,11 +133,11 @@ def insert_argrole(edge: Hyperedge, argrole: str, pos: int) -> Hyperedge:
         return replace_argroles(edge, argroles)
     else:
         st = edge.mtype()
-        if st in {"C", "R"}:
+        if st in {EdgeType.CONCEPT, EdgeType.RELATION}:
             new_edge = [insert_argrole(edge[0], argrole, pos)]
             new_edge += edge[1:]
             return Hyperedge(new_edge)
-        elif st in {"P", "B"}:
+        elif st in {EdgeType.PREDICATE, EdgeType.BUILDER}:
             new_edge = [edge[0], insert_argrole(edge[1], argrole, pos)]
             new_edge += list(edge[2:])
             return Hyperedge(new_edge)
