code organization

Author: emmettbutler

ddtrace/contrib/dd_trace_api/__init__.py

@@ -1,6 +1,8 @@
 from sys import addaudithook
+from typing import Any
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Tuple
 import weakref
 
@@ -32,25 +34,43 @@ def _proxy_span_arguments(args: List, kwargs: Dict) -> Tuple[List, Dict]:
     return proxied_args, proxied_kwargs
 
 
-def _patched(method_of, fn_name):
-    def _inner(state_shared_with_api, *args, **kwargs):
-        retval_from_api = state_shared_with_api.get("api_return_value")
-        operand_stub = state_shared_with_api.get("stub_self")
-        args, kwargs = _proxy_span_arguments(args, kwargs)
-        retval_from_impl = getattr(_STUB_TO_REAL[operand_stub], fn_name)(*args, **kwargs)
-        if "impl_return_value" in state_shared_with_api:
-            state_shared_with_api["impl_return_value"] = retval_from_impl
-        if retval_from_api is not None:
-            _STUB_TO_REAL[retval_from_api] = retval_from_impl
-
-    return _inner
+def _call_on_real_instance(
+    operand_stub: dd_trace_api._Stub,
+    method_name: str,
+    retval_from_api: Optional[Any],
+    state_shared_with_api: Dict,
+    *args: List,
+    **kwargs: Dict
+) -> None:
+    """
+    Call `method_name` on the real object corresponding to `operand_stub` with `args` and `kwargs` as arguments.
+    Pass the return value back to the API layer via the mutable `state_shared_with_api`.
+
+    Store the value that will be returned from the API call we're in the middle of, for the purpose
+    of mapping from those Stub objects to their real counterparts.
+    """
+    args, kwargs = _proxy_span_arguments(args, kwargs)
+    retval_from_impl = getattr(_STUB_TO_REAL[operand_stub], method_name)(*args, **kwargs)
+    if "impl_return_value" in state_shared_with_api:
+        state_shared_with_api["impl_return_value"] = retval_from_impl
+    if retval_from_api is not None:
+        _STUB_TO_REAL[retval_from_api] = retval_from_impl
 
 
 def _hook(name, hook_args):
+    """Called in response to `sys.audit` events"""
     if not dd_trace_api.__datadog_patch or not name.startswith(_DD_HOOK_PREFIX):
         return
     args = hook_args[0][0]
-    _patched(*(name.replace(_DD_HOOK_PREFIX, "").split(".")))(args[0], *args[1:], **hook_args[0][1])
+    state_shared_with_api = args[0]
+    _call_on_real_instance(
+        state_shared_with_api.get("stub_self"),
+        name.replace(_DD_HOOK_PREFIX, "").rsplit(".", 1)[-1],
+        state_shared_with_api.get("api_return_value"),
+        state_shared_with_api,
+        *args[1:],
+        **hook_args[0][1]
+    )
 
 
 def get_version() -> str: