Repository: Instagram/MonkeyType

Issue:
Calls to locally defined functions are not captured
Calls to locally defined functions are not captured, for example:

`mod.py`:
```python
def add(a, b):
    def add_impl(a, b):
        return a + b
    return add_impl(a, b)
```
`main.py`:
```python
from mod import add
add(1, 2)
```

After running `monkeytype run ./main.py`, the resulting `monkeytype.sqlite3` contains just the call to `add`.

Discussion:
carljm: This looks reasonable as far as it goes. Thanks!

But what is the impact on stub generation of having these traces in the trace store? Is stub generation able to handle them reasonably? How about `monkeytype apply`?

Patch:
diff --git a/monkeytype/tracing.py b/monkeytype/tracing.py
index 4109b6b..1360d08 100644
--- a/monkeytype/tracing.py
+++ b/monkeytype/tracing.py
@@ -130,6 +130,17 @@ def _has_code(
     return None
 
 
+def get_previous_frames(frame: Optional[FrameType]) -> Iterator[FrameType]:
+    while frame is not None:
+        yield frame
+        frame = frame.f_back
+
+
+def get_locals_from_previous_frames(frame: FrameType) -> Iterator[Any]:
+    for previous_frame in get_previous_frames(frame):
+        yield from previous_frame.f_locals.values()
+
+
 def get_func(frame: FrameType) -> Optional[Callable[..., Any]]:
     """Return the function whose code object corresponds to the supplied stack frame."""
     code = frame.f_code
@@ -154,6 +165,14 @@ def get_func(frame: FrameType) -> Optional[Callable[..., Any]]:
             func = get_func_in_mro(v, code)
             if func is not None:
                 break
+    # If we still can't find the function, try looking at the locals of all previous frames.
+    if func is None:
+        for v in get_locals_from_previous_frames(frame):
+            if not callable(v):
+                continue
+            func = _has_code(v, code)
+            if func is not None:
+                break
     return func
 
 
@@ -215,7 +234,7 @@ class CallTracer:
         if frame in self.traces:
             # resuming a generator; we've already seen this frame
             return
-        arg_names = code.co_varnames[0 : code.co_argcount]
+        arg_names = code.co_varnames[: code.co_argcount + code.co_kwonlyargcount]
         arg_types = {}
         for name in arg_names:
             if name in frame.f_locals:


Test patch:
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
index 04b39f9..efd4552 100644
--- a/tests/test_tracing.py
+++ b/tests/test_tracing.py
@@ -6,8 +6,10 @@
 import inspect
 from types import FrameType
 from typing import (
+    Callable,
     Iterator,
     Optional,
+    Tuple,
 )
 
 import pytest
@@ -39,8 +41,8 @@ def simple_add(a: int, b: int) -> int:
     return a + b
 
 
-def has_optional_kwarg(a: int, b: str = None) -> Optional[FrameType]:
-    return inspect.currentframe()
+def uses_kw_only_arg(a: int, *, b: int) -> int:
+    return a + b
 
 
 def has_locals(foo: str) -> Optional[FrameType]:
@@ -128,6 +130,21 @@ async def sum_squares(n: int) -> int:
     return tot
 
 
+def call_method_on_locally_defined_class(n: int) -> Tuple[type, Callable]:
+    class Math:
+        def square(self, n: int) -> int:
+            return n * n
+    Math().square(n)
+    return Math, Math.square
+
+
+def call_locally_defined_function(n: int) -> Callable:
+    def square(n: int) -> int:
+        return n * n
+    square(n)
+    return square
+
+
 def implicit_return_none() -> None:
     pass
 
@@ -187,6 +204,11 @@ class TestTraceCalls:
             simple_add(1, 2)
         assert collector.traces == [CallTrace(simple_add, {'a': int, 'b': int}, int)]
 
+    def test_kw_only_arg(self, collector):
+        with trace_calls(collector, max_typed_dict_size=0):
+            uses_kw_only_arg(1, b=2)
+        assert collector.traces == [CallTrace(uses_kw_only_arg, {'a': int, 'b': int}, int)]
+
     def test_flushes(self, collector):
         with trace_calls(collector, max_typed_dict_size=0):
             pass
@@ -242,6 +264,18 @@ class TestTraceCalls:
                 pass
         assert collector.traces == [CallTrace(squares, {'n': int}, NoneType, int)]
 
+    def test_locally_defined_class_trace(self, collector):
+        with trace_calls(collector, max_typed_dict_size=0):
+            cls, method = call_method_on_locally_defined_class(3)
+        assert len(collector.traces) == 2
+        assert collector.traces[0] == CallTrace(method, {'self': cls, 'n': int}, int)
+
+    def test_locally_defined_function_trace(self, collector):
+        with trace_calls(collector, max_typed_dict_size=0):
+            function = call_locally_defined_function(3)
+        assert len(collector.traces) == 2
+        assert collector.traces[0] == CallTrace(function, {'n': int}, int)
+
     def test_return_none(self, collector):
         """Ensure traces have a return_type of NoneType for functions that return a value of None"""
         with trace_calls(collector, max_typed_dict_size=0):