Improves docstrings for overloaded functions

Improves the docstrings for overloaded functions to be stylistically similar
to non-overloaded functions.

Also updates the helpers that generate strings from type annotations to be
more consistent with the style the documentation uses. For example, `Union[int, float]`
will now be rendered as `int | float`.

Author: pranavm-nvidia

tripy/docs/_static/style.css

@@ -21,3 +21,14 @@ section {
     margin-top: 2rem;
     margin-bottom: 2rem;
 }
+
+.func-overload-sig {
+    padding-left: 3em !important;
+    color: var(--color-api-overall);
+    font-style: normal;
+}
+
+.func-overload-sig p {
+    margin-bottom: 0 !important;
+    margin-top: 0 !important;
+}

tripy/tests/test_function_registry.py

@@ -25,7 +25,7 @@
 
 import tripy as tp
 from tripy import TripyException
-from tripy.function_registry import AnnotationInfo, FunctionRegistry, render_arg_type, sanitize_name
+from tripy.function_registry import AnnotationInfo, FunctionRegistry, type_str_from_arg, str_from_type_annotation
 
 
 @pytest.fixture()
@@ -199,10 +199,10 @@ def func(a: int):
 
         func_overload = registry.overloads["test"][0]
 
-        assert not func_overload.annotations
+        assert not func_overload._annotations
         assert registry["test"](0) == 1
-        assert func_overload.annotations
-        assert func_overload.annotations["a"] == AnnotationInfo(int, False, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        assert func_overload._annotations
+        assert func_overload._annotations["a"] == AnnotationInfo(int, False, inspect.Parameter.POSITIONAL_OR_KEYWORD)
 
     def test_doc_of_non_overloaded_func(self, registry):
         # When there is no overload, the registry function should
@@ -224,31 +224,46 @@ def func(a: int):
             """
             pass
 
+        # Tripy types should turn into class links
         @registry("test")
-        def func(a: float):
+        def func(a: Union[int, "tripy.Tensor"]):
             """
-            This func takes a float.
+            This func takes an int or a tensor.
             """
             pass
 
         print(registry["test"].__doc__)
         assert (
             registry["test"].__doc__
             == dedent(
-                """
+                r"""
                 *This function has multiple overloads:*
 
                 ----------
 
-                > **test** (*a*: :class:`int`) -> None
+                .. role:: sig-prename
+                    :class: sig-prename descclassname
+                .. role:: sig-name
+                    :class: sig-name descname
+
+                .. container:: func-overload-sig sig sig-object py
+
+                    :sig-prename:`tripy`\ .\ :sig-name:`test`\ (a: int) -> None
 
                 This func takes an int.
 
                 ----------
 
-                > **test** (*a*: :class:`float`) -> None
+                .. role:: sig-prename
+                    :class: sig-prename descclassname
+                .. role:: sig-name
+                    :class: sig-name descname
+
+                .. container:: func-overload-sig sig sig-object py
+
+                    :sig-prename:`tripy`\ .\ :sig-name:`test`\ (a: int | :class:`tripy.Tensor`) -> None
 
-                This func takes a float.
+                This func takes an int or a tensor.
                 """
             ).strip()
         )
@@ -379,7 +394,7 @@ def func(n: Union[int, float]) -> int:
                   [0-9]+ \|     \.\.\.
                       \|\s
 
-                Not a valid overload because: For parameter: 'n', expected an instance of type: 'Union\[int, float\]' but got argument of type: 'List\[str\]'\.
+                Not a valid overload because: For parameter: 'n', expected an instance of type: 'int | float' but got argument of type: 'List\[str\]'\.
             """
             ).strip(),
         ):
@@ -403,7 +418,7 @@ def func(n: Sequence[int]) -> int:
                   [0-9]+ \|     \.\.\.
                       \|\s
 
-                Not a valid overload because: For parameter: 'n', expected an instance of type: 'Sequence\[int\]' but got argument of type: 'List\[Union\[(int, str)|(str, int)\]\]'\.
+                Not a valid overload because: For parameter: 'n', expected an instance of type: 'Sequence\[int\]' but got argument of type: 'List\[(int \| str)|(str \| int)\]'\.
             """
             ).strip(),
         ):
@@ -475,7 +490,7 @@ def func(n: Sequence[Union[int, float]]) -> int:
                   [0-9]+ \|     \.\.\.
                       \|\s
 
-                Not a valid overload because: For parameter: 'n', expected an instance of type: 'Sequence\[Union\[int, float\]\]' but got argument of type: 'List\[str\]'\.
+                Not a valid overload because: For parameter: 'n', expected an instance of type: 'Sequence\[int | float\]' but got argument of type: 'List\[str\]'\.
             """
             ).strip(),
         ):
@@ -496,16 +511,16 @@ def func(a: int, *args: int) -> int:
 @pytest.mark.parametrize(
     "typ, expected",
     [
-        (tp.types.TensorLike, "Union[tripy.Tensor, numbers.Number]"),
-        (tp.types.ShapeLike, "Sequence[Union[int, tripy.DimensionSize]]"),
+        (tp.types.TensorLike, "tripy.Tensor | numbers.Number"),
+        (tp.types.ShapeLike, "Sequence[int | tripy.DimensionSize]"),
         (tp.Tensor, "Tensor"),
         (torch.Tensor, "torch.Tensor"),
         (int, "int"),
-        (Optional[int], "Optional[int]"),
+        (Optional[int], "int | None"),
     ],
 )
-def test_sanitize_name(typ, expected):
-    assert sanitize_name(typ) == expected
+def test_str_from_type_annotation(typ, expected):
+    assert str_from_type_annotation(typ) == expected
 
 
 @pytest.mark.parametrize(
@@ -517,5 +532,5 @@ def test_sanitize_name(typ, expected):
         ("hi", "str"),
     ],
 )
-def test_render_arg_type(typ, expected):
-    assert render_arg_type(typ) == expected
+def test_type_str_from_arg(typ, expected):
+    assert type_str_from_arg(typ) == expected

tripy/tripy/backend/api/executable.py

@@ -23,7 +23,7 @@
 from tripy.backend.mlir import utils as mlir_utils
 from tripy.common.exception import raise_error
 from tripy.frontend import Tensor
-from tripy.function_registry import sanitize_name
+from tripy.function_registry import str_from_type_annotation
 from tripy.utils import json as json_utils
 from dataclasses import dataclass
 
@@ -73,8 +73,11 @@ def stream(self, stream):
         self._executor.stream = stream
 
     def __str__(self) -> str:
-        params = [f"{name}: {sanitize_name(param.annotation)}" for name, param in self.__signature__.parameters.items()]
-        return f"Executable({', '.join(params)}) -> {sanitize_name(self.__signature__.return_annotation)}"
+        params = [
+            f"{name}: {str_from_type_annotation(param.annotation)}"
+            for name, param in self.__signature__.parameters.items()
+        ]
+        return f"Executable({', '.join(params)}) -> {str_from_type_annotation(self.__signature__.return_annotation)}"
 
     @staticmethod
     def load(path: str) -> "tripy.Executable":

tripy/tripy/frontend/ops/tensor_initializers.py

@@ -292,7 +292,7 @@ def arange(
 ) -> "tripy.Tensor":
     r"""
     Returns a 1D tensor containing a sequence of numbers in the half-open interval
-    :math:`[0, \text{stop})` incrementing by :math:`\text{step}`.
+    :math:`[\text{start}, \text{stop})` incrementing by :math:`\text{step}`.
 
     Args:
         start: The inclusive lower bound of the values to generate. If a tensor is provided, it must be a scalar tensor.