Microoptimize the walk() function

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: