Add get_present_trait method to statement class (#370)

Noticed in QuEraComputing/bloqade-circuit#163
that this is useful for linting, but also just for shortening code here
and there. I also replaced code that raises a `ValueError` if the result
from `get_trait` is `None` here (2nd commit).

Author: david-pl

src/kirin/ir/method.py

@@ -68,24 +68,18 @@ def arg_types(self):
     @property
     def self_type(self):
         """Return the type of the self argument of the method."""
-        trait = self.code.get_trait(HasSignature)
-        if trait is None:
-            raise ValueError("Method body must implement HasSignature")
+        trait = self.code.get_present_trait(HasSignature)
         signature = trait.get_signature(self.code)
         return Generic(Method, Generic(tuple, *signature.inputs), signature.output)
 
     @property
     def callable_region(self):
-        trait = self.code.get_trait(CallableStmtInterface)
-        if trait is None:
-            raise ValueError("Method body must implement CallableStmtInterface")
+        trait = self.code.get_present_trait(CallableStmtInterface)
         return trait.get_callable_region(self.code)
 
     @property
     def return_type(self):
-        trait = self.code.get_trait(HasSignature)
-        if trait is None:
-            raise ValueError("Method body must implement HasSignature")
+        trait = self.code.get_present_trait(HasSignature)
         return trait.get_signature(self.code).output
 
     def __repr__(self) -> str:

src/kirin/ir/nodes/stmt.py

@@ -689,6 +689,15 @@ def get_trait(cls, trait: type[TraitType]) -> TraitType | None:
                 return t
         return None
 
+    @classmethod
+    def get_present_trait(cls, trait: type[TraitType]) -> TraitType:
+        """Just like get_trait, but expects the trait to be there.
+        Useful for linter checks, when you know the trait is present."""
+        for t in cls.traits:
+            if isinstance(t, trait):
+                return t
+        raise ValueError(f"Trait {trait} not present in statement {cls}")
+
     def expect_one_result(self) -> ResultValue:
         """Check if the statement contain only one result, and return it"""
         if len(self._results) != 1: