Improvement: Docstring Parsing for Method Description (#518)

* tmp impl * add docstring parser to branch * pysdk upgrade related * just for sphinx * minor, cleanup * minor, cleanup * add tests for different docstring formatting options * add changelog for docstrings * rm prints * Update pyteal/ast/subroutine.py Co-authored-by: Hang Su <[email protected]> * Update pyteal/ast/subroutine.py Co-authored-by: Hang Su <[email protected]> * Apply suggestions from code review Co-authored-by: Hang Su <[email protected]> * add strict version requirement, better testing, address double newline for longer descriptions * fmt * add comments * add version pin on docstring parser in setup.py * linted * maybe fix description * Use new docstring features in AlgoBank (#520) * dont add return descr if it returns void * unneceary change (#522) * fix newlines in arg descs * fix newline * Update pyteal/ast/subroutine.py Co-authored-by: Hang Su <[email protected]> * line end spaces Co-authored-by: Hang Su <[email protected]> Co-authored-by: Hang Su <[email protected]> Co-authored-by: Jason Paulos <[email protected]>
algorand · Sep 6, 2022 · 59d6aa0 · 59d6aa0
1 parent 632c94e
commit 59d6aa0
Show file tree

Hide file tree

Showing 7 changed files with 261 additions and 34 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Unreleased
 
+## Added
+
+* ABI Methods will now parse the docstring for the method and set the description for any parameters that are described. ([#518](https://github.com/algorand/pyteal/pull/518))
+  * Note: the docstring must adhere to one of google, rst, numpy , or epy formatting styles.
+
 ## Fixed
 * Subroutines annotated with a `TupleX` class are now invoked with an instance of that exact class, instead of the more general `Tuple` class ([#519](https://github.com/algorand/pyteal/pull/519))
 

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -3,3 +3,4 @@ sphinx-rtd-theme==1.0.0
 # dependencies from setup.py
 py-algorand-sdk>=1.9.0,<2.0.0
 semantic-version>=2.9.0,<3.0.0
+docstring-parser==0.14.1
diff --git a/examples/application/abi/algobank.json b/examples/application/abi/algobank.json
@@ -6,28 +6,32 @@
             "args": [
                 {
                     "type": "pay",
-                    "name": "payment"
+                    "name": "payment",
+                    "desc": "A payment transaction containing the amount of Algos the user wishes to deposit. The receiver of this transaction must be this app's escrow account."
                 },
                 {
                     "type": "account",
-                    "name": "sender"
+                    "name": "sender",
+                    "desc": "An account that is opted into this app (or will opt in during this method call). The deposited funds will be recorded in this account's local state. This account must be the same as the sender of the `payment` transaction."
                 }
             ],
             "returns": {
                 "type": "void"
             },
-            "desc": "This method receives a payment from an account opted into this app and records it in their local state. The caller may opt into this app during this call."
+            "desc": "This method receives a payment from an account opted into this app and records it as a deposit.\nThe caller may opt into this app during this call."
         },
         {
             "name": "getBalance",
             "args": [
                 {
                     "type": "account",
-                    "name": "user"
+                    "name": "user",
+                    "desc": "The user whose balance you wish to look up. This user must be opted into this app."
                 }
             ],
             "returns": {
-                "type": "uint64"
+                "type": "uint64",
+                "desc": "The balance corresponding to the given user, in microAlgos."
             },
             "desc": "Lookup the balance of a user held by this app."
         },
@@ -36,17 +40,19 @@
             "args": [
                 {
                     "type": "uint64",
-                    "name": "amount"
+                    "name": "amount",
+                    "desc": "The amount of Algos requested to be withdraw, in microAlgos. This method will fail if this amount exceeds the amount of Algos held by this app for the method call sender."
                 },
                 {
                     "type": "account",
-                    "name": "recipient"
+                    "name": "recipient",
+                    "desc": "An account who will receive the withdrawn Algos. This may or may not be the same as the method call sender."
                 }
             ],
             "returns": {
                 "type": "void"
             },
-            "desc": "Withdraw an amount of Algos held by this app. The sender of this method call will be the source of the Algos, and the destination will be the `recipient` argument. This may or may not be the same as the sender's address. This method will fail if the amount of Algos requested to be withdrawn exceeds the amount of Algos held by this app for the sender. The Algos will be transferred to the recipient using an inner transaction whose fee is set to 0, meaning the caller's transaction must include a surplus fee to cover the inner transaction."
+            "desc": "Withdraw an amount of Algos held by this app.\nThe sender of this method call will be the source of the Algos, and the destination will be the `recipient` argument.\nThe Algos will be transferred to the recipient using an inner transaction whose fee is set to 0, meaning the caller's transaction must include a surplus fee to cover the inner transaction."
         }
     ],
     "networks": {}

diff --git a/examples/application/abi/algobank.py b/examples/application/abi/algobank.py
@@ -41,10 +41,16 @@ def assert_sender_is_creator() -> Expr:
 
 @router.method(no_op=CallConfig.CALL, opt_in=CallConfig.CALL)
 def deposit(payment: abi.PaymentTransaction, sender: abi.Account) -> Expr:
-    """This method receives a payment from an account opted into this app and records it in
-    their local state.
+    """This method receives a payment from an account opted into this app and records it as a deposit.
 
     The caller may opt into this app during this call.
+
+    Args:
+        payment: A payment transaction containing the amount of Algos the user wishes to deposit.
+            The receiver of this transaction must be this app's escrow account.
+        sender: An account that is opted into this app (or will opt in during this method call).
+            The deposited funds will be recorded in this account's local state. This account must
+            be the same as the sender of the `payment` transaction.
     """
     return Seq(
         Assert(payment.get().sender() == sender.address()),
@@ -59,7 +65,14 @@ def deposit(payment: abi.PaymentTransaction, sender: abi.Account) -> Expr:
 
 @router.method
 def getBalance(user: abi.Account, *, output: abi.Uint64) -> Expr:
-    """Lookup the balance of a user held by this app."""
+    """Lookup the balance of a user held by this app.
+
+    Args:
+        user: The user whose balance you wish to look up. This user must be opted into this app.
+
+    Returns:
+        The balance corresponding to the given user, in microAlgos.
+    """
     return output.set(App.localGet(user.address(), Bytes("balance")))
 
 
@@ -68,14 +81,17 @@ def withdraw(amount: abi.Uint64, recipient: abi.Account) -> Expr:
     """Withdraw an amount of Algos held by this app.
 
     The sender of this method call will be the source of the Algos, and the destination will be
-    the `recipient` argument. This may or may not be the same as the sender's address.
-
-    This method will fail if the amount of Algos requested to be withdrawn exceeds the amount of
-    Algos held by this app for the sender.
+    the `recipient` argument.
 
     The Algos will be transferred to the recipient using an inner transaction whose fee is set
     to 0, meaning the caller's transaction must include a surplus fee to cover the inner
     transaction.
+
+    Args:
+        amount: The amount of Algos requested to be withdraw, in microAlgos. This method will fail
+            if this amount exceeds the amount of Algos held by this app for the method call sender.
+        recipient: An account who will receive the withdrawn Algos. This may or may not be the same
+            as the method call sender.
     """
     return Seq(
         # if amount is larger than App.localGet(Txn.sender(), Bytes("balance")), the subtraction

diff --git a/pyteal/ast/subroutine.py b/pyteal/ast/subroutine.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from docstring_parser import parse as parse_docstring
 from inspect import isclass, Parameter, signature, get_annotations
 from types import MappingProxyType, NoneType
 from typing import Any, Callable, Final, Optional, TYPE_CHECKING, cast
@@ -610,32 +611,73 @@ def method_signature(self, overriding_name: str = None) -> str:
         return f"{overriding_name}({','.join(args)}){self.type_of()}"
 
     def method_spec(self) -> sdk_abi.Method:
-        skip_names = ["return", "output"]
+        desc: str = ""
+        arg_descs: dict[str, str] = {}
+        return_desc: str = ""
+        args: list[dict[str, str]] = []
 
-        args = [
-            {
-                "type": str(abi.type_spec_from_annotation(val)),
-                "name": name,
-            }
-            for name, val in self.subroutine.annotations.items()
-            if name not in skip_names
-        ]
+        if self.subroutine.implementation.__doc__:
+            docstring = parse_docstring(self.subroutine.implementation.__doc__)
 
-        spec = {
-            "name": self.name(),
-            "args": args,
-            "returns": {"type": str(self.type_of())},
-        }
+            # Combine short and long descriptions with newline
+            method_descriptions: list[str] = []
+            if docstring.short_description:
+                method_descriptions.append(docstring.short_description)
+            if docstring.long_description:
+                method_descriptions.append(docstring.long_description)
 
-        if self.subroutine.implementation.__doc__ is not None:
-            spec["desc"] = " ".join(
+            method_desc = "\n\n".join(method_descriptions)
+
+            # Turn double new line into single, replacing single newline with space
+            desc = "\n".join(
                 [
-                    i.strip()
-                    for i in self.subroutine.implementation.__doc__.split("\n")
-                    if not (i.isspace() or len(i) == 0)
+                    i.replace("\n", " ").strip()
+                    for i in method_desc.split("\n\n")
+                    if i.strip()
                 ]
             )
 
+            # Get the descriptions for any documented arguments
+            arg_descs = {
+                arg.arg_name: arg.description.replace("\n", " ").strip()
+                for arg in docstring.params
+                if arg.arg_name != self.OUTPUT_ARG_NAME and arg.description is not None
+            }
+
+            # Get the special return description
+            return_desc = (
+                ""
+                if not docstring.returns or not docstring.returns.description
+                else docstring.returns.description.replace("\n", " ").strip()
+            )
+
+            # Generate the ABI method object given the subroutine args
+            # Add in description if one is set
+            for name, val in self.subroutine.annotations.items():
+                # Skip annotations for `return` and `output` in the args list
+                if name in ["return", self.OUTPUT_ARG_NAME]:
+                    continue
+
+                arg_obj = {
+                    "type": str(abi.type_spec_from_annotation(val)),
+                    "name": name,
+                }
+
+                if name in arg_descs:
+                    arg_obj["desc"] = arg_descs[name]
+
+                args.append(arg_obj)
+
+        # Create the return obj for the method, adding description if set
+        return_obj = {"type": str(self.type_of())}
+        if return_desc and return_obj["type"] != "void":
+            return_obj["desc"] = return_desc
+
+        # Create the method spec, adding description if set
+        spec = {"name": self.name(), "args": args, "returns": return_obj}
+        if desc:
+            spec["desc"] = desc
+
         return sdk_abi.Method.undictify(spec)
 
     def type_of(self) -> str | abi.TypeSpec:

diff --git a/pyteal/ast/subroutine_test.py b/pyteal/ast/subroutine_test.py
@@ -1302,3 +1302,159 @@ def mySubroutine(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
         actual, _ = declaration.__teal__(options)
         options.setSubroutine(None)
         assert actual == expected
+
+
+def test_docstring_parsing_with_different_format():
+    short_desc = "Example of a ABIReturnSubroutine with short description docstring."
+    a_doc = "an abi Uint64 value"
+    return_doc = "A PyTeal expression that sets output Uint64 value as argument a."
+    long_desc = """Example first line.
+
+    This is a second line.
+
+    This is a third line that's so long it has to wrap in order to fit properly
+    in a line of source code.
+    """
+    expected_long_desc = "Example first line.\nThis is a second line.\nThis is a third line that's so long it has to wrap in order to fit properly in a line of source code."
+
+    def documented_method(a: pt.abi.Uint64, *, output: pt.abi.Uint64):
+        return output.set(a)
+
+    # Google format
+    documented_method.__doc__ = f"""{short_desc}
+
+    Args:
+        a: {a_doc}
+
+    Returns:
+        {return_doc}
+    """
+
+    mspec_dict = pt.ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == short_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # epy format
+    documented_method.__doc__ = f"""
+    {short_desc}
+
+    @param a: {a_doc}
+    @return: {return_doc}
+    """
+
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == short_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # numpy format
+    documented_method.__doc__ = f"""{short_desc}
+
+    Parameters
+    ----------
+    a:
+        an abi Uint64 value
+    output:
+        {a_doc}
+
+    Returns
+    -------
+    uint64
+        {return_doc}
+    """
+
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == short_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # rst format
+    documented_method.__doc__ = f"""{short_desc}
+
+    :param a: {a_doc}
+    :returns: {return_doc}
+    """
+
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == short_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # Short and long descriptions
+    documented_method.__doc__ = f"""{long_desc}
+
+    :param a: {a_doc}
+    :returns: {return_doc}
+    """
+
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == expected_long_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # Only long description
+    # Short description is defined as being on the first line, so by introducing
+    # long_desc on the second, there is no short description.
+    documented_method.__doc__ = f"""
+    {long_desc}
+
+    :param a: {a_doc}
+    :returns: {return_doc}
+    """
+
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert mspec_dict["desc"] == expected_long_desc
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    # No description
+    documented_method.__doc__ = f"""
+
+    :param a: {a_doc}
+    :returns: {return_doc}
+    """
+    mspec_dict = ABIReturnSubroutine(documented_method).method_spec().dictify()
+    assert "descr" not in mspec_dict
+    assert mspec_dict["args"][0]["desc"] == a_doc
+    assert mspec_dict["returns"]["desc"] == return_doc
+
+    algobank_example = """Withdraw an amount of Algos held by this app.
+
+    The sender of this method call will be the source of the Algos, and the destination will be
+    the `recipient` argument.
+
+    The Algos will be transferred to the recipient using an inner transaction whose fee is set
+    to 0, meaning the caller's transaction must include a surplus fee to cover the inner
+    transaction.
+
+    Args:
+        amount: The amount of Algos requested to be withdraw, in microAlgos. This method will fail
+            if this amount exceeds the amount of Algos held by this app for the method call sender.
+        recipient: An account who will receive the withdrawn Algos. This may or may not be the same
+            as the method call sender.
+    """
+
+    # algobank example
+    def withdraw(amount: pt.abi.Uint64, recipient: pt.abi.Account):
+        return pt.Assert(pt.Int(1))
+
+    withdraw.__doc__ = algobank_example
+
+    mspec_dict = ABIReturnSubroutine(withdraw).method_spec().dictify()
+    assert (
+        mspec_dict["desc"]
+        == "Withdraw an amount of Algos held by this app.\nThe sender of this method call will be the source of the Algos, "
+        + "and the destination will be the `recipient` argument.\nThe Algos will be transferred to the recipient using an inner transaction whose fee is "
+        + "set to 0, meaning the caller's transaction must include a surplus fee to cover the inner transaction."
+    )
+    assert (
+        mspec_dict["args"][0]["desc"]
+        == "The amount of Algos requested to be withdraw, in microAlgos. This method will fail if this amount exceeds "
+        + "the amount of Algos held by this app for the method call sender."
+    )
+    assert (
+        mspec_dict["args"][1]["desc"]
+        == "An account who will receive the withdrawn Algos. This may or may not be the same as the method call sender."
+    )
+    assert "desc" not in mspec_dict["returns"]