Use __set_name__ for better wrapping.

anntzer · anntzer · commit 98534b60ad93 · 2020-04-05T00:45:38.000+02:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1,6 +1,5 @@
 from collections import OrderedDict
 from contextlib import ExitStack
-import functools
 import inspect
 import itertools
 import logging
@@ -31,42 +30,59 @@
 _log = logging.getLogger(__name__)
 
 
-def _axis_method_wrapper(attr_name, method_name, *, doc_sub=None):
+class _axis_method_wrapper:
     """
     Helper to generate Axes methods wrapping Axis methods.
 
     After ::
 
         get_foo = _axis_method_wrapper("xaxis", "get_bar")
 
-    ``get_foo`` is a method that forwards it arguments to the ``get_bar``
-    method of the ``xaxis`` attribute, and gets its signature and docstring
-    from ``Axis.get_bar``.
+    (in the body of a class) ``get_foo`` is a method that forwards it arguments
+    to the ``get_bar`` method of the ``xaxis`` attribute, and gets its
+    signature and docstring from ``Axis.get_bar``.
 
     The docstring of ``get_foo`` is built by replacing "this Axis" by "the
     {attr_name}" ("the xaxis", "the yaxis") in the wrapped method's docstring;
     additional replacements can by given in *doc_sub*.  The docstring is also
     dedented to simplify further manipulations.
     """
 
-    method = getattr(maxis.Axis, method_name)
-    get_method = attrgetter(f"{attr_name}.{method_name}")
-
-    @functools.wraps(method)
-    def wrapper(self, *args, **kwargs):
-        return get_method(self)(*args, **kwargs)
-
-    doc = wrapper.__doc__
-    if doc:
-        doc_sub = {"this Axis": f"the {attr_name}", **(doc_sub or {})}
-        for k, v in doc_sub.items():
-            assert k in doc, \
-                (f"The docstring of wrapped Axis method {method_name!r} must "
-                 f"contain {k!r} as a substring.")
-            doc = doc.replace(k, v)
-        wrapper.__doc__ = inspect.cleandoc(doc)
-
-    return wrapper
+    def __init__(self, attr_name, method_name, *, doc_sub=None):
+        self.attr_name = attr_name
+        self.method_name = method_name
+        self.doc_sub = doc_sub
+
+    def __set_name__(self, owner, name):
+        # This is called at the end of the class body as
+        # ``self.__set_name__(cls, name_under_which_self_is_assigned)``; we
+        # rely on that to give the wrapper the correct __name__/__qualname__.
+        get_method = attrgetter(f"{self.attr_name}.{self.method_name}")
+
+        def wrapper(self, *args, **kwargs):
+            return get_method(self)(*args, **kwargs)
+
+        wrapper.__module__ = owner.__module__
+        wrapper.__name__ = name
+        wrapper.__qualname__ = f"{owner.__qualname__}.{name}"
+        # Manually copy the signature instead of using functools.wraps because
+        # displaying the Axis method source when asking for the Axes method
+        # source would be confusing.
+        wrapped_method = getattr(maxis.Axis, self.method_name)
+        wrapper.__signature__ = inspect.signature(wrapped_method)
+        doc = wrapped_method.__doc__
+        if doc:
+            doc_sub = {"this Axis": f"the {self.attr_name}",
+                       **(self.doc_sub or {})}
+            for k, v in doc_sub.items():
+                assert k in doc, \
+                    (f"The definition of {wrapper.__qualname__} expected that "
+                     f"the docstring of Axis.{self.method_name} contains "
+                     f"{k!r} as a substring.")
+                doc = doc.replace(k, v)
+            wrapper.__doc__ = inspect.cleandoc(doc)
+
+        setattr(owner, name, wrapper)
 
 
 def _process_plot_format(fmt):
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -829,7 +829,7 @@ def set_zscale(self, value, **kwargs):
     get_zticklabels = _axis_method_wrapper("zaxis", "get_ticklabels")
     set_zticklabels = _axis_method_wrapper(
         "zaxis", "_set_ticklabels",
-        doc_sub={"Axis.set_ticks": "Axes.set_zticks"})
+        doc_sub={"Axis.set_ticks": "Axes3D.set_zticks"})
 
     zaxis_date = _axis_method_wrapper("zaxis", "axis_date")
     if zaxis_date.__doc__:
