..

python/sdist/amici/swig.py

@@ -1,4 +1,5 @@
 """Functions related to SWIG or SWIG-generated code"""
+from __future__ import annotations
 import ast
 import contextlib
 import re
@@ -74,7 +75,10 @@ def visit_FunctionDef(self, node):
                 arg.annotation = self._new_annot(arg.annotation.value)
         return node
 
-    def _new_annot(self, old_annot: str):
+    def _new_annot(self, old_annot: str | ast.Name):
+        if isinstance(old_annot, ast.Name):
+            old_annot = old_annot.id
+
         with contextlib.suppress(KeyError):
             return self.mapping[old_annot]