Add docstring with examples for array_api (#574)

* add docstring to creation functions

* reorder elementwise_functions

* run doctests

* doctest_modules -> doctest-modules

* doctest_modules -> doctest-modules

* add env variables

* remove env. variables

* add docstring for statistical_functions

* update datatypes exported

* add docstring for datatypes

* add examples to most elementwise functions

* add more docstrings

* temporarely disable doctests on CI

* move step

.github/workflows/rbc_test.yml

@@ -7,10 +7,6 @@ on:
     branches:
       - main
 
-env:
-  NUMBA_LOOP_VECTORIZE: 0
-  NUMBA_SLP_VECTORIZE: 0
-
 
 # kill any previous running job on a new commit
 concurrency:
@@ -231,8 +227,6 @@ jobs:
           RBC_TESTS_FULL: TRUE
         run: |
           mamba run -n rbc pytest -sv -r A rbc/tests/heavydb/ -x
-          pkill -f heavydb
-          mamba run -n rbc pytest -sv -r A rbc/tests/ -x --ignore rbc/tests/heavydb
 
       - name: Run rbc tests
         shell: bash -l {0}
@@ -245,8 +239,13 @@ jobs:
           RBC_TESTS_FULL: TRUE
         run: |
           mamba run -n rbc pytest -sv -r A rbc/tests/heavydb/ -x
-          pkill -f heavydb
-          mamba run -n rbc pytest -sv -r A rbc/tests/ -x --ignore rbc/tests/heavydb
+
+      - name: Run doctests
+        shell: bash -l {0}
+        # requires HeavyDB 6.4 or newer
+        if: ${{ false }}
+        run: |
+          mamba run -n rbc pytest -sv -r A rbc/stdlib/ --doctest-modules -x
 
       - name: Show Heavydb conda logs on failure [conda]
         shell: bash -l {0}
@@ -271,6 +270,17 @@ jobs:
           mamba run -n heavydb-env killall heavydb
           sleep 5
 
+      - name: Run Non-HeavyDB tests
+        shell: bash -l {0}
+        env:
+          HEAVYDB_SOURCE: ${{ matrix.heavydb-from }}
+          EXPECTED_PYTHON_VERSION: ${{ matrix.python-version }}
+          EXPECTED_HEAVYDB_VERSION: ${{ matrix.heavydb-version }}
+          EXPECTED_NUMBA_VERSION: ${{ matrix.numba-version }}
+          RBC_TESTS_FULL: TRUE
+        run: |
+          mamba run -n rbc pytest -sv -r A rbc/tests/ -x --ignore rbc/tests/heavydb
+
       - name: Show Heavydb server output [conda]
         shell: bash -l {0}
         if: matrix.os == 'ubuntu-latest' && matrix.heavydb-from == 'conda'

doc/conf.py

@@ -94,3 +94,13 @@
     "github_version": "main",
     "doc_path": "doc",
 }
+
+
+# -- ignore certain parts of the docstring -----------------------------------
+from sphinx.ext.autodoc import between
+
+def setup(app):
+    # Register a sphinx.ext.autodoc.between listener to ignore everything
+    # between lines that contain the word IGNORE
+    app.connect('autodoc-process-docstring', between('^.*IGNORE.*$', exclude=True))
+    return app

rbc/stdlib/__init__.py

@@ -1,13 +1,12 @@
+import textwrap
 import functools
 from enum import Enum
 from numba.core import extending
 from rbc.heavydb import Array, ArrayPointer
 from rbc import typesystem, errors
 
-
-ARRAY_API_ADDRESS = ("https://data-apis.org/array-api/latest/API_specification"
-                     "/generated/signatures.{0}.{1}.html"
-                     "#signatures.{0}.{1}")
+ARRAY_API_ADDRESS = ("https://data-apis.org/array-api/latest/API_specification/"
+                     "generated/array_api.{0}.html#array_api.{0}")
 NUMPY_API_ADDRESS = ("https://numpy.org/doc/stable/reference/generated/numpy.{0}.html")
 ADDRESS = ARRAY_API_ADDRESS
 
@@ -37,6 +36,7 @@ def determine_input_type(argty):
 class Expose:
     def __init__(self, globals, module_name):
         self._globals = globals
+        # XXX: remove module_name as it is not used
         self.module_name = module_name
 
     def create_function(self, func_name):
@@ -56,10 +56,10 @@ def format_docstring(self, ov_func, func_name, api):
             # Array API
             link = (
                 f"`Array-API '{func_name}' "
-                f"doc <{ARRAY_API_ADDRESS.format(self.module_name, func_name)}>`_")
+                f"doc <{ARRAY_API_ADDRESS.format(func_name)}>`_")
 
         if original_docstring is not None:
-            new_doctring = f"{original_docstring}\n\n{link}"
+            new_doctring = f"{textwrap.dedent(original_docstring)}\n\n{link}"
         else:
             new_doctring = link
         return new_doctring
@@ -95,12 +95,12 @@ def wraps(overload_func):
 
 class BinaryUfuncExpose(Expose):
 
-    def implements(self, ufunc, ufunc_name=None, dtype=None, api=API.ARRAY_API):
+    def implements(self, ufunc, func_name=None, dtype=None, api=API.ARRAY_API):
         """
         Wrapper for binary ufuncs that returns an array
         """
-        if ufunc_name is None:
-            ufunc_name = ufunc.__name__
+        if func_name is None:
+            func_name = ufunc.__name__
 
         def binary_ufunc_impl(a, b):
             typA = determine_input_type(a)
@@ -154,10 +154,10 @@ def impl(a, b):
                     return nb_dtype(ufunc(cast_a, cast_b))
                 return impl
 
-        fn = self.create_function(ufunc_name)
+        fn = self.create_function(func_name)
 
         def wrapper(overload_func):
-            overload_func.__doc__ = self.format_docstring(overload_func, ufunc_name, api)
+            overload_func.__doc__ = self.format_docstring(overload_func, func_name, api)
             functools.update_wrapper(fn, overload_func)
 
             decorate = extending.overload(fn)
@@ -168,12 +168,12 @@ def wrapper(overload_func):
 
 class UnaryUfuncExpose(BinaryUfuncExpose):
 
-    def implements(self, ufunc, ufunc_name=None, dtype=None, api=API.ARRAY_API):
+    def implements(self, ufunc, func_name=None, dtype=None, api=API.ARRAY_API):
         """
         Wrapper for unary ufuncs that returns an array
         """
-        if ufunc_name is None:
-            ufunc_name = ufunc.__name__
+        if func_name is None:
+            func_name = ufunc.__name__
 
         def unary_ufunc_impl(a):
             nb_dtype = determine_dtype(a, dtype)
@@ -196,10 +196,10 @@ def impl(a):
                     return nb_dtype(ufunc(cast))
                 return impl
 
-        fn = self.create_function(ufunc_name)
+        fn = self.create_function(func_name)
 
         def wrapper(overload_func):
-            overload_func.__doc__ = self.format_docstring(overload_func, ufunc_name, api)
+            overload_func.__doc__ = self.format_docstring(overload_func, func_name, api)
             functools.update_wrapper(fn, overload_func)
 
             decorate = extending.overload(fn)

rbc/stdlib/creation_functions.py

@@ -108,6 +108,28 @@ def _array_api_triu(x, k=0):
 def _impl__full(shape, fill_value, dtype=None):
     """
     Return a new array of given shape and type, filled with fill_value.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('double[](int64, double)')
+    ... def rbc_full(size, fill_value):
+    ...     return array_api.full(size, fill_value)
+    >>> rbc_full(3, 3.14).execute()
+    array([3.14, 3.14, 3.14], dtype=float32)
+
+    >>> @heavydb('int64[](int64, int64)')
+    ... def rbc_full2(size, fill_value):
+    ...     return array_api.full(size, fill_value, 'int64')
+    >>> rbc_full2(2, 0).execute()
+    array([0, 0])
     """
     nb_dtype = _determine_dtype(dtype, fill_value)
 
@@ -122,6 +144,28 @@ def impl(shape, fill_value, dtype=None):
 def _impl_full_like(a, fill_value, dtype=None):
     """
     Return a full array with the same shape and type as a given array.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('double[](double[], double)')
+    ... def full_like(other, fill_value):
+    ...     return array_api.full_like(other, fill_value)
+    >>> full_like([1.0, 2.0, 3.0], 3.14).execute()
+    array([3.14, 3.14, 3.14], dtype=float32)
+
+    >>> @heavydb('int64[](int64[], int64)')
+    ... def full_like(other, fill_value):
+    ...     return array_api.full_like(other, fill_value, 'int64')
+    >>> full_like([1, 2, 3], 3).execute()
+    array([3, 3, 3])
     """
     if isinstance(a, ArrayPointer):
         if dtype is None:
@@ -175,6 +219,28 @@ def impl(shape, dtype=None):
 def _impl_zeros(shape, dtype=None):
     """
     Return a new array of given shape and type, filled with zeros.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('double[](int64)')
+    ... def zeros(size):
+    ...     return array_api.zeros(size)
+    >>> zeros(3).execute()
+    array([0., 0., 0.], dtype=float32)
+
+    >>> @heavydb('int64[](int64)')
+    ... def zeros2(size):
+    ...     return array_api.zeros(size, 'int64')
+    >>> zeros2(2).execute()
+    array([0, 0])
     """
 
     nb_dtype = _determine_dtype(dtype, fill_value=None)
@@ -189,6 +255,28 @@ def impl(shape, dtype=None):
 def _impl_zeros_like(a, dtype=None):
     """
     Return an array of zeros with the same shape and type as a given array.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('int64[](int64[])')
+    ... def zeros_like(other):
+    ...     return array_api.zeros_like(other)
+    >>> zeros_like([1, 2, 3]).execute()
+    array([0, 0, 0])
+
+    >>> @heavydb('float[](int64[])')
+    ... def zeros_like2(other):
+    ...     return array_api.zeros_like(other, 'float32')
+    >>> zeros_like2([1, 2, 3]).execute()
+    array([0., 0., 0.], dtype=float32)
     """
     if isinstance(a, ArrayPointer):
         if dtype is None:
@@ -207,6 +295,28 @@ def impl(a, dtype=None):
 def _impl_ones(shape, dtype=None):
     """
     Return a new array of given shape and type, filled with ones.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('double[](int64)')
+    ... def ones(size):
+    ...     return array_api.ones(size)
+    >>> ones(3).execute()
+    array([1., 1., 1.], dtype=float32)
+
+    >>> @heavydb('int64[](int64)')
+    ... def ones2(size):
+    ...     return array_api.ones(size, 'int64')
+    >>> ones2(2).execute()
+    array([1, 1])
     """
 
     nb_dtype = _determine_dtype(dtype, fill_value=None)
@@ -221,6 +331,28 @@ def impl(shape, dtype=None):
 def _impl_ones_like(a, dtype=None):
     """
     Return an array of ones with the same shape and type as a given array.
+
+    Examples
+    --------
+    IGNORE:
+    >>> from rbc.heavydb import global_heavydb_singleton
+    >>> heavydb = next(global_heavydb_singleton)
+    >>> heavydb.unregister()
+
+    IGNORE
+
+    >>> from rbc.stdlib import array_api
+    >>> @heavydb('int64[](int64[])')
+    ... def ones_like(other):
+    ...     return array_api.ones_like(other)
+    >>> ones_like([1, 2, 3]).execute()
+    array([1, 1, 1])
+
+    >>> @heavydb('float[](int64[])')
+    ... def ones_like2(other):
+    ...     return array_api.ones_like(other, 'float32')
+    >>> ones_like2([1, 2, 3]).execute()
+    array([1., 1., 1.], dtype=float32)
     """
     if isinstance(a, ArrayPointer):
         if dtype is None: