Expand record return values for all function calls

Author: aktech

record_api/core.py

@@ -29,6 +29,13 @@
 context_manager: Optional[ContextManager] = None
 write_line: Optional[Callable[[dict], None]] = None
 
+FUNCTION_CALL_OP_NAMES = {
+    "CALL_METHOD",
+    "CALL_FUNCTION",
+    "CALL_FUNCTION_KW",
+    "CALL_FUNCTION_EX",
+}
+
 
 def get_tracer() -> Tracer:
     global TRACER
@@ -400,7 +407,8 @@ def __call__(self) -> None:
             return self.process(
                 (self.TOS, self.TOS1), BINARY_OPS[opname], (self.TOS1, self.TOS)
             )
-        if self.previous_stack and self.previous_stack.opname == "CALL_METHOD":
+
+        if self.previous_stack and self.previous_stack.opname in FUNCTION_CALL_OP_NAMES:
             self.log_called_method()
 
         method_name = f"op_{opname}"
@@ -409,15 +417,16 @@ def __call__(self) -> None:
         return None
 
     def log_called_method(self):
-        filename, line, fn, args, *kwargs = self.previous_stack.log_call_args
-        kwargs = kwargs[0] if kwargs else {}
-        log_call(
-            f"{filename}:{line}",
-            fn,
-            tuple(args),
-            *((kwargs,) if kwargs else ()),
-            return_type=type(self.TOS),
-        )
+        if self.previous_stack.log_call_args:
+            filename, line, fn, args, *kwargs = self.previous_stack.log_call_args
+            kwargs = kwargs[0] if kwargs else {}
+            log_call(
+                f"{filename}:{line}",
+                fn,
+                tuple(args),
+                *((kwargs,) if kwargs else ()),
+                return_type=type(self.TOS),
+            )
 
     # special case subscr b/c we only check first arg, not both
     def op_BINARY_SUBSCR(self):
@@ -489,7 +498,7 @@ def op_COMPARE_OP(self):
     def op_CALL_FUNCTION(self):
         args = self.pop_n(self.oparg)
         fn = self.pop()
-        self.process((fn,), fn, args)
+        self.process((fn,), fn, args, delay=True)
 
     def op_CALL_FUNCTION_KW(self):
         kwargs_keys = self.pop()
@@ -499,7 +508,7 @@ def op_CALL_FUNCTION_KW(self):
         args = self.pop_n(self.oparg - n_kwargs)
         fn = self.pop()
 
-        self.process((fn,), fn, args, kwargs)
+        self.process((fn,), fn, args, kwargs, delay=True)
 
     def op_CALL_FUNCTION_EX(self):
         has_kwarg = self.oparg & int("01", 2)
@@ -513,7 +522,7 @@ def op_CALL_FUNCTION_EX(self):
             fn = self.pop()
         if inspect.isgenerator(args):
             return
-        self.process((fn,), fn, args, kwargs)
+        self.process((fn,), fn, args, kwargs, delay=True)
 
     def op_CALL_METHOD(self):
         args = self.pop_n(self.oparg)

record_api/test.py

@@ -139,20 +139,20 @@ def test_sort(self):
         self.trace("self.a.sort(axis=0)")
         self.assertCalls(
             call(ANY, getattr, (self.a, "sort")),
-            call(ANY, self.a.sort, (), {"axis": 0}),
+            call(ANY, self.a.sort, (), {"axis": 0}, return_type=type(None)),
         )
 
     def test_eye(self):
         self.trace("np.eye(10, order='F')")
         self.assertCalls(
-            call(ANY, getattr, (np, "eye")), call(ANY, np.eye, (10,), {"order": "F"}),
+            call(ANY, getattr, (np, "eye")), call(ANY, np.eye, (10,), {"order": "F"}, return_type=np.ndarray),
         )
 
     def test_linspace(self):
         self.trace("np.linspace(3, 4, endpoint=False)")
         self.assertCalls(
             call(ANY, getattr, (np, "linspace",)),
-            call(ANY, np.linspace, (3, 4,), {"endpoint": False}),
+            call(ANY, np.linspace, (3, 4,), {"endpoint": False}, return_type=np.ndarray),
         )
 
     def test_reshape(self):
@@ -167,7 +167,7 @@ def test_concatenate(self):
         self.trace("np.concatenate((self.a, self.a), axis=0)")
         self.assertCalls(
             call(ANY, getattr, (np, "concatenate",)),
-            call(ANY, np.concatenate, ((self.a, self.a),), {"axis": 0}),
+            call(ANY, np.concatenate, ((self.a, self.a),), {"axis": 0}, return_type=np.ndarray),
         )
 
     def test_ravel_list(self):
@@ -196,7 +196,7 @@ def test_numpy_array_constructor(self):
         self.trace("np.ndarray(dtype='int64', shape=tuple())")
         self.assertCalls(
             call(ANY, getattr, (np, "ndarray")),
-            call(ANY, np.ndarray, (), {"dtype": "int64", "shape": tuple()}),
+            call(ANY, np.ndarray, (), {"dtype": "int64", "shape": tuple()}, return_type=np.ndarray),
         )
 
     def test_not_contains(self):