..

python/sdist/amici/swig.py

@@ -68,7 +68,7 @@ def visit_FunctionDef(self, node):
             for arg in node.args.args:
                 if not arg.annotation:
                     continue
-                if isinstance(arg.annotation, ast.Name):
+                if not isinstance(arg.annotation, ast.Constant):
                     # there is already proper annotation
                     continue
 
@@ -133,15 +133,15 @@ def _annotation_from_docstring(self, node: ast.FunctionDef):
         for line_no, line in enumerate(docstring):
             if (
                 match := re.match(
-                    r"\s*:rtype:\s*(?::py:class:`)?(\w+)`?\s*$", line
+                    r"\s*:rtype:\s*(?::py:class:`)?(.+)`?\s*$", line
                 )
             ) and not match.group(1).startswith(":"):
                 node.returns = ast.Constant(match.group(1))
                 lines_to_remove.add(line_no)
 
             if (
                 match := re.match(
-                    r"\s*:type\s*(\w+):\W*(?::py:class:`)?(\w+)`?\s*$", line
+                    r"\s*:type\s*(\w+):\W*(?::py:class:`)?(.+)`?\s*$", line
                 )
             ) and not match.group(1).startswith(":"):
                 for arg in node.args.args:

swig/amici.i

@@ -373,13 +373,8 @@ __all__ = [
 
 
 def _get_ptr(
-    obj: Union[AmiciModel, AmiciExpData, AmiciSolver, AmiciReturnData],
-) -> Union[
-    Model,
-    ExpData,
-    Solver,
-    ReturnData,
-]:
+    obj: AmiciModel | AmiciExpData | AmiciSolver | AmiciReturnData,
+) -> Model | ExpData | Solver | ReturnData:
     """
     Convenience wrapper that returns the smart pointer pointee, if applicable