Upgrading remove_returns and prepend_docstr_noreturns to remove_sections and prepend_docstr_nosections.

bsipocz · bsipocz · commit e97067f050ee · 2017-10-24T13:04:45.000+01:00
The new functions remove and prepend a list of sections rather than just “Returns”
diff --git a/astroquery/utils/__init__.py b/astroquery/utils/__init__.py
@@ -10,4 +10,4 @@
 from .class_or_instance import *
 from .commons import *
 from .process_asyncs import async_to_sync
-from .docstr_chompers import prepend_docstr_noreturns
+from .docstr_chompers import prepend_docstr_nosections
diff --git a/astroquery/utils/docstr_chompers.py b/astroquery/utils/docstr_chompers.py
@@ -10,22 +10,22 @@ def dec(fn):
     return dec
 
 
-def prepend_docstr_noreturns(doc):
+def prepend_docstr_nosections(doc, sections=['Returns', ]):
     """
     Decorator to prepend to the function's docstr after stripping out the
-    "Returns".
+    list of sections provided (by default "Returns" only).
     """
     def dec(fn):
-        fn.__doc__ = ("\n".join(remove_returns(doc)) +
+        fn.__doc__ = ("\n".join(remove_sections(doc, sections)) +
                       textwrap.dedent(fn.__doc__))
         return fn
     return dec
 
 
-def remove_returns(doc):
+def remove_sections(doc, sections):
     """
-    Given a numpy-formatted docstring, remove the "Returns" block
-    and dedent the whole thing.
+    Given a numpy-formatted docstring, remove the section blocks provided in
+    ``sections`` and dedent the whole thing.
 
     Returns
     -------
@@ -37,7 +37,7 @@ def remove_returns(doc):
     rblock = False
     for line in lines:
         lstrip = line.rstrip()
-        if lstrip == "Returns":
+        if lstrip in sections:
             rblock = True
             continue
         elif rblock:
diff --git a/astroquery/utils/tests/test_utils.py b/astroquery/utils/tests/test_utils.py
@@ -19,7 +19,7 @@
 
 from ...utils import chunk_read, chunk_report, class_or_instance, commons
 from ...utils.process_asyncs import async_to_sync_docstr, async_to_sync
-from ...utils.docstr_chompers import remove_returns, prepend_docstr_noreturns
+from ...utils.docstr_chompers import remove_sections, prepend_docstr_nosections
 
 
 class SimpleQueryClass(object):
@@ -284,6 +284,10 @@ def test_async_to_sync(cls=Dummy):
 
 
 docstr3 = """
+    Parameters
+    ----------
+    first_param
+
     Returns
     -------
     Nothing!
@@ -301,10 +305,12 @@ def test_async_to_sync(cls=Dummy):
 
 
 def test_return_chomper(doc=docstr3, out=docstr3_out):
-    assert remove_returns(doc) == [x.lstrip() for x in out.split('\n')]
+    assert (remove_sections(doc, sections=['Returns', 'Parameters']) ==
+                [x.lstrip() for x in out.split('\n')])
+
 
 
-def dummyfunc():
+def dummyfunc1():
     """
     Returns
     -------
@@ -317,15 +323,28 @@ def dummyfunc():
     pass
 
 
+def dummyfunc2():
+    """
+    Returns
+    -------
+    Nothing!
+    """
+    pass
+
+
 docstr4 = """
     Blah Blah Blah
 
     Returns
     -------
     nothing
+
+    Examples
+    --------
+    no_examples_at_all
 """
 
-docstr4_out = """
+docstr4_out1 = """
     Blah Blah Blah
 
     Returns
@@ -337,10 +356,20 @@ def dummyfunc():
     Nada
 """
 
+docstr4_out2 = """
+    Blah Blah Blah
+
+    Returns
+    -------
+    Nothing!
+"""
+
 
-def test_prepend_docstr(doc=docstr4, func=dummyfunc, out=docstr4_out):
-    fn = prepend_docstr_noreturns(doc)(func)
-    assert fn.__doc__ == textwrap.dedent(docstr4_out)
+@pytest.mark.parametrize("func, out", [(dummyfunc1, docstr4_out1),
+                                       (dummyfunc2, docstr4_out2)])
+def test_prepend_docstr(func, out, doc=docstr4):
+    fn = prepend_docstr_nosections(doc, sections=['Returns', 'Examples'])(func)
+    assert fn.__doc__ == textwrap.dedent(out)
 
 
 @async_to_sync
