Skip to content

Commit

Permalink
Add and test PLUM_SIMPLE_DOC
Browse files Browse the repository at this point in the history
  • Loading branch information
wesselb committed Sep 15, 2023
1 parent 57e11e6 commit 2ccaff4
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 0 deletions.
6 changes: 6 additions & 0 deletions plum/function.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import textwrap
from functools import wraps
from types import MethodType
Expand Down Expand Up @@ -141,6 +142,11 @@ def __doc__(self) -> Optional[str]:
# clearing the cache.
self.clear_cache(reregister=False)

# Don't do any fancy appending of docstrings when the environment variable
# `PLUM_SIMPLE_DOC` is set to `1`.
if "PLUM_SIMPLE_DOC" in os.environ and os.environ["PLUM_SIMPLE_DOC"] == "1":
return self._doc

# Derive the basis of the docstring from `self._f`, removing any indentation.
doc = self._doc.strip()
if doc:
Expand Down
26 changes: 26 additions & 0 deletions tests/test_function.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import abc
import os
import textwrap
import typing

Expand Down Expand Up @@ -221,6 +222,31 @@ def f(x: int):
assert g.__doc__ == textwrap.dedent(expected_doc).strip()


def test_simple_doc(monkeypatch):
@dispatch
def f(x: int):
"""First."""

@dispatch
def f(x: str):
"""Second."""

monkeypatch.setitem(os.environ, "PLUM_SIMPLE_DOC", "1")
assert f.__doc__ == "First."

monkeypatch.setitem(os.environ, "PLUM_SIMPLE_DOC", "0")
expected_doc = """
First.
-----------
f(x: str)
Second.
"""
assert f.__doc__ == textwrap.dedent(expected_doc).strip()


def test_methods():
dispatch = Dispatcher()

Expand Down

0 comments on commit 2ccaff4

Please sign in to comment.