💫 Improve introspection of custom extension attributes (#3729)

* Add custom __dir__ to Underscore (see #3707)

* Make sure custom extension methods keep their docstrings (see #3707)

* Improve tests

* Prepend note on partial to docstring (see #3707)

* Remove print statement

* Handle cases where docstring is None

spacy/tests/doc/test_underscore.py

@@ -140,3 +140,28 @@ def test_underscore_mutable_defaults_dict(en_vocab):
     assert len(token1._.mutable) == 2
     assert token1._.mutable["x"] == ["y"]
     assert len(token2._.mutable) == 0
+
+
+def test_underscore_dir(en_vocab):
+    """Test that dir() correctly returns extension attributes. This enables
+    things like tab-completion for the attributes in doc._."""
+    Doc.set_extension("test_dir", default=None)
+    doc = Doc(en_vocab, words=["hello", "world"])
+    assert "_" in dir(doc)
+    assert "test_dir" in dir(doc._)
+    assert "test_dir" not in dir(doc[0]._)
+    assert "test_dir" not in dir(doc[0:2]._)
+
+
+def test_underscore_docstring(en_vocab):
+    """Test that docstrings are available for extension methods, even though
+    they're partials."""
+
+    def test_method(doc, arg1=1, arg2=2):
+        """I am a docstring"""
+        return (arg1, arg2)
+
+    Doc.set_extension("test_docstrings", method=test_method)
+    doc = Doc(en_vocab, words=["hello", "world"])
+    assert test_method.__doc__ == "I am a docstring"
+    assert doc._.test_docstrings.__doc__.rsplit(". ")[-1] == "I am a docstring"

spacy/tokens/underscore.py

@@ -25,14 +25,28 @@ def __init__(self, extensions, obj, start=None, end=None):
         object.__setattr__(self, "_start", start)
         object.__setattr__(self, "_end", end)
 
+    def __dir__(self):
+        # Hack to enable autocomplete on custom extensions
+        extensions = list(self._extensions.keys())
+        return ["set", "get", "has"] + extensions
+
     def __getattr__(self, name):
         if name not in self._extensions:
             raise AttributeError(Errors.E046.format(name=name))
         default, method, getter, setter = self._extensions[name]
         if getter is not None:
             return getter(self._obj)
         elif method is not None:
-            return functools.partial(method, self._obj)
+            method_partial = functools.partial(method, self._obj)
+            # Hack to port over docstrings of the original function
+            # See https://stackoverflow.com/q/27362727/6400719
+            method_docstring = method.__doc__ or ""
+            method_docstring_prefix = (
+                "This method is a partial function and its first argument "
+                "(the object it's called on) will be filled automatically. "
+            )
+            method_partial.__doc__ = method_docstring_prefix + method_docstring
+            return method_partial
         else:
             key = self._get_key(name)
             if key in self._doc.user_data: