Fix upcasting by using annotated assignment

opshin/prelude.py

@@ -62,8 +62,11 @@ def resolve_spent_utxo(txins: List[TxInInfo], p: Spending) -> TxOut:
     return [txi.resolved for txi in txins if txi.out_ref == p.tx_out_ref][0]
 
 
-def resolve_datum(txout: TxOut, tx_info: TxInfo) -> BuiltinData:
-    """Returns the datum attached to a given transaction output, independent of whether it was inlined or embedded."""
+def resolve_datum_unsafe(txout: TxOut, tx_info: TxInfo) -> BuiltinData:
+    """
+    Returns the datum attached to a given transaction output, independent of whether it was inlined or embedded.
+    Raises an exception if no datum was attached.
+    """
     attached_datum = txout.datum
     if isinstance(attached_datum, SomeOutputDatumHash):
         # TODO can we raise a KeyError here if not found?
@@ -77,3 +80,23 @@ def resolve_datum(txout: TxOut, tx_info: TxInfo) -> BuiltinData:
         # no datum attached
         assert False, "No datum was attached to the given transaction output"
     return res
+
+
+def resolve_datum(
+    txout: TxOut, tx_info: TxInfo
+) -> Union[SomeOutputDatum, NoOutputDatum]:
+    """
+    Returns SomeOutputDatum with the datum attached to a given transaction output,
+    independent of whether it was inlined or embedded, if there was an attached datum.
+    Otherwise it returns NoOutputDatum.
+    """
+    attached_datum = txout.datum
+    res: Union[SomeOutputDatum, NoOutputDatum] = NoOutputDatum()
+    if isinstance(attached_datum, SomeOutputDatumHash):
+        if attached_datum.datum_hash in tx_info.data.keys():
+            res = SomeOutputDatum(
+                tx_info.data.get(attached_datum.datum_hash, Nothing())
+            )
+    else:
+        res = attached_datum
+    return res

opshin/rewrite/rewrite_import_hashlib.py

@@ -23,6 +23,9 @@ def attribute(self, attr) -> plt.AST:
             return plt.Lambda(["self"], plt.Var("self"))
         raise NotImplementedError("HashType only has attribute 'digest'")
 
+    def __ge__(self, other):
+        return isinstance(other, HashType)
+
 
 HashInstanceType = InstanceType(HashType())
 

opshin/type_inference.py

@@ -63,7 +63,10 @@ def exit_scope(self):
         self.scopes.pop()
 
     def set_variable_type(self, name: str, typ: Type, force=False):
-        if not force and name in self.scopes[-1] and typ != self.scopes[-1][name]:
+        if not force and name in self.scopes[-1] and self.scopes[-1][name] != typ:
+            if self.scopes[-1][name] >= typ:
+                # the specified type is broader, we pass on this
+                return
             raise TypeInferenceError(
                 f"Type {self.scopes[-1][name]} of variable {name} in local scope does not match inferred type {typ}"
             )
@@ -209,9 +212,10 @@ def visit_AnnAssign(self, node: AnnAssign) -> TypedAnnAssign:
             node.target.id, InstanceType(typed_ass.annotation), force=True
         )
         typed_ass.target = self.visit(node.target)
-        assert typed_ass.value.typ >= InstanceType(
-            typed_ass.annotation
-        ), "Can only downcast to a specialized type"
+        assert (
+            typed_ass.value.typ >= InstanceType(typed_ass.annotation)
+            or InstanceType(typed_ass.annotation) >= typed_ass.value.typ
+        ), "Can only cast between related types"
         return typed_ass
 
     def visit_If(self, node: If) -> TypedIf: