diff --git a/spacy/tests/doc/test_underscore.py b/spacy/tests/doc/test_underscore.py index 8f47157fa..2877bfeea 100644 --- a/spacy/tests/doc/test_underscore.py +++ b/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" diff --git a/spacy/tokens/underscore.py b/spacy/tokens/underscore.py index ef1d78717..b36fe9294 100644 --- a/spacy/tokens/underscore.py +++ b/spacy/tokens/underscore.py @@ -25,6 +25,11 @@ class Underscore(object): 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)) @@ -32,7 +37,16 @@ class Underscore(object): 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: