libcore: pass return type of method from stub to proxy via type annotation

Author: cnvogelg

amitools/vamos/libcore/proxy.py

@@ -1,3 +1,5 @@
+from typing import get_type_hints
+
 from amitools.vamos.machine import Code, REG_D0, REG_D1
 from amitools.vamos.libtypes import TagList
 
@@ -18,11 +20,12 @@ def __init__(self, ctx, base_addr=None, run_sp=None):
 
 
 class LibProxyRegs:
-    def __init__(self, ctx, args, arg_regs, kwargs):
+    def __init__(self, ctx, args, arg_regs, kwargs, stub_method=None):
         assert len(args) == len(arg_regs)
         self.ctx = ctx
         self.args = args
         self.arg_regs = arg_regs
+        self.stub_method = stub_method
         # shall we return d1 as well?
         self.ret_d1 = kwargs.pop("ret_d1", False)
         # shall we wrap the return value into a type
@@ -83,7 +86,14 @@ def return_d1(self):
         return self.ret_d1
 
     def wrap_result(self):
-        return self.wrap_res
+        # either use forced type in proxy call argument
+        result_type = self.wrap_res
+        # or use from type hint of stub method (if any)
+        if not result_type and self.stub_method:
+            hints = get_type_hints(self.stub_method)
+            if hints and "return" in hints:
+                result_type = hints["return"]
+        return result_type
 
     def cleanup(self):
         for mem in self.auto_strings:
@@ -109,7 +119,7 @@ def _gen_arg_regs(self, func_def):
     def _gen_stub_call(self, arg_regs, stub_method):
         def stub_call(self, *args, **kwargs):
             """proxy function to call lib stub directly"""
-            regs = LibProxyRegs(self.ctx, args, arg_regs, kwargs)
+            regs = LibProxyRegs(self.ctx, args, arg_regs, kwargs, stub_method)
 
             # fill registers with arg values
             # (lib call may depend on it)

amitools/vamos/libcore/stub.py

@@ -1,4 +1,5 @@
 from types import MethodType
+from typing import get_type_hints
 import time
 import traceback
 
@@ -358,6 +359,19 @@ def profile_func(this, *args, **kwargs):
 
         return profile_func
 
+    def _copy_return_type_hint(self, in_func, out_func):
+        # get type hint
+        hints = get_type_hints(in_func)
+
+        # is a return type defined
+        if hints and "return" in hints:
+            return_type = hints["return"]
+            # attach new return annotation
+            ann = out_func.__annotations__
+            if ann is None:
+                ann = out_func.__annotations__ = {}
+            ann["return"] = return_type
+
     def _wrap_func(self, stub, impl_func, ctx, profile):
         """create a stub func for a valid impl func
         returns an unbound method for the stub instaance
@@ -389,4 +403,7 @@ def _wrap_func(self, stub, impl_func, ctx, profile):
         if profile:
             func = self._gen_profile_func(fd_func, profile, func)
 
+        # copy return hint from method to wrap func for later use in proxy
+        self._copy_return_type_hint(method, func)
+
         return func