perf: optimize the walk() function #587

I noticed that it takes pynvim about 4ms to attach to an nvim instance for me,
and 3ms of that is due to the single line:
  metadata = walk(decode_if_bytes, metadata)

This commit reduces the walk() time down to 1.5ms, which brings the total
attach time down to 2.5ms. This is helpful for me because in my use case I end
up connecting to all of the currently-running nvim processes and this starts to
take a noticeable amount of time. Unfortunately parallelization does not help
here due to the nature of the slowness.

walk() is expensive because it does a very large amount of pure-python
manipulation, so this commit is just some tweaks to reduce the overheads:

- *args and **kw make the function call slow, and we can avoid needing them by
  pre-packing the args into fn via functools.partial
- The comprehensions can be written to directly construct the objects rather
  than create a generator which is passed to a constructor
- The typechecking is microoptimized by calling type() once and unrolling the
  `type_ in [list, tuple]` check

I did notice that in my setup the metadata contains no byte objects, so the
entire call is a noop. I'm not sure if that is something that could be relied
on or detected, which could be an even bigger speedup.

Author: kmod

pynvim/api/common.py

@@ -241,11 +241,15 @@ def decode_if_bytes(obj: T, mode: TDecodeMode = True) -> Union[T, str]:
     return obj
 
 
-def walk(fn: Callable[..., Any], obj: Any, *args: Any, **kwargs: Any) -> Any:
-    """Recursively walk an object graph applying `fn`/`args` to objects."""
-    if type(obj) in [list, tuple]:
-        return list(walk(fn, o, *args) for o in obj)
-    if type(obj) is dict:
-        return dict((walk(fn, k, *args), walk(fn, v, *args)) for k, v in
-                    obj.items())
-    return fn(obj, *args, **kwargs)
+def walk(fn: Callable[[Any], Any], obj: Any) -> Any:
+    """Recursively walk an object graph applying `fn` to objects."""
+
+    # Note: this function is very hot, so it is worth being careful
+    # about performance.
+    type_ = type(obj)
+
+    if type_ is list or type_ is tuple:
+        return [walk(fn, o) for o in obj]
+    if type_ is dict:
+        return {walk(fn, k): walk(fn, v) for k, v in obj.items()}
+    return fn(obj)

pynvim/api/nvim.py

@@ -209,7 +209,7 @@ def request(self, name: str, *args: Any, **kwargs: Any) -> Any:
         decode = kwargs.pop('decode', self._decode)
         args = walk(self._to_nvim, args)
         res = self._session.request(name, *args, **kwargs)
-        return walk(self._from_nvim, res, decode=decode)
+        return walk(partial(self._from_nvim, decode=decode), res)
 
     def next_message(self) -> Any:
         """Block until a message(request or notification) is available.

pynvim/plugin/host.py

@@ -112,7 +112,7 @@ def _wrap_delayed_function(self, cls, delayed_handlers, name, sync,
 
     def _wrap_function(self, fn, sync, decode, nvim_bind, name, *args):
         if decode:
-            args = walk(decode_if_bytes, args, decode)
+            args = walk(partial(decode_if_bytes, mode=decode), args)
         if nvim_bind is not None:
             args.insert(0, nvim_bind)
         try: