Add --doctest-ufunc option to doctest Numpy ufuncs

CHANGES.rst

@@ -1,6 +1,8 @@
 0.12.0 (unreleased)
 ===================
 
+- Run doctests in docstrings of Numpy ufuncs. [#123, #174]
+
 0.11.2 (2021-12-09)
 ===================
 

README.rst

@@ -42,6 +42,7 @@ providing the following features:
 * handling doctests that use remote data in conjunction with the
   `pytest-remotedata`_ plugin (see `Remote Data`_)
 * optional inclusion of ``*.rst`` files for doctests (see `Setup and Configuration`_)
+* inclusion of doctests in docstrings of Numpy ufuncs
 
 .. _pytest-remotedata: https://github.com/astropy/pytest-remotedata
 

pytest_doctestplus/plugin.py

@@ -157,6 +157,20 @@ def get_optionflags(parent):
     return flag_int
 
 
+def _is_numpy_ufunc(method):
+    try:
+        import numpy as np
+    except ModuleNotFoundError:
+        # If Numpy is not installed, then there can't be any ufuncs!
+        return False
+    while True:
+        try:
+            method = method.__wrapped__
+        except AttributeError:
+            break
+    return isinstance(method, np.ufunc)
+
+
 def pytest_configure(config):
     doctest_plugin = config.pluginmanager.getplugin('doctest')
     run_regular_doctest = config.option.doctestmodules and not config.option.doctest_plus
@@ -238,7 +252,13 @@ def collect(self):
             runner = doctest.DebugRunner(
                 verbose=False, optionflags=options, checker=OutputChecker())
 
-            for test in finder.find(module):
+            tests = finder.find(module)
+            for method in module.__dict__.values():
+                if _is_numpy_ufunc(method):
+                    found = finder.find(method, module=module)
+                    tests += found
+
+            for test in tests:
                 if test.examples:  # skip empty doctests
                     ignore_warnings_context_needed = False
                     show_warnings_context_needed = False

setup.cfg

@@ -38,6 +38,7 @@ install_requires =
 
 [options.extras_require]
 test =
+    numpy
     pytest-remotedata>=0.3.2
     sphinx
 

tests/test_doctestplus.py

@@ -1,6 +1,8 @@
+import glob
 import os
 from packaging.version import Version
 from textwrap import dedent
+import sys
 
 import pytest
 
@@ -930,3 +932,105 @@ def test_remote_data_ignore_warnings(testdir):
     )
     testdir.inline_run(p, '--doctest-plus', '--doctest-rst', '--remote-data').assertoutcome(passed=1)
     testdir.inline_run(p, '--doctest-plus', '--doctest-rst').assertoutcome(skipped=1)
+
+
+def test_ufunc(testdir):
+    pytest.importorskip('numpy')
+
+    # Create and build example module
+    testdir.makepyfile(module1="""
+        def foo():
+            '''A doctest...
+
+            >>> foo()
+            1
+            '''
+            return 1
+        """)
+    testdir.makepyfile(module2="""
+        from _module2 import foo
+        """)
+    testdir.makepyfile(setup="""
+        from setuptools import setup, Extension
+        import numpy as np
+
+        ext = Extension('_module2', ['_module2.c'],
+                        extra_compile_args=['-std=c99'],
+                        include_dirs=[np.get_include()])
+        setup(name='example', py_modules=['module1', 'module2'], ext_modules=[ext])
+        """)
+    testdir.makefile('.c', _module2=r"""
+        #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+
+        #include <numpy/arrayobject.h>
+        #include <numpy/ufuncobject.h>
+        #include <Python.h>
+
+
+        static double foo_inner(double a, double b)
+        {
+            return a + b;
+        }
+
+
+        static void foo_loop(
+            char **args,
+            const npy_intp *dimensions,
+            const npy_intp *steps,
+            void *NPY_UNUSED(data)
+        ) {
+            const npy_intp n = dimensions[0];
+            for (npy_intp i = 0; i < n; i ++)
+            {
+                *(double *) &args[2][i * steps[2]] = foo_inner(
+                *(double *) &args[0][i * steps[0]],
+                *(double *) &args[1][i * steps[1]]);
+            }
+        }
+
+
+        static PyUFuncGenericFunction foo_loops[] = {foo_loop};
+        static char foo_types[] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE};
+        static void *foo_data[] = {NULL};
+        static const char foo_name[] = "foo";
+        static const char foo_docstring[] = ">>> foo(1, 2)\n3.0";
+
+        static PyModuleDef moduledef = {
+            .m_base = PyModuleDef_HEAD_INIT,
+            .m_name = "_module2",
+            .m_size = -1
+        };
+
+
+        PyMODINIT_FUNC PyInit__module2(void)
+        {
+            import_array();
+            import_ufunc();
+
+            PyObject *module = PyModule_Create(&moduledef);
+            if (!module)
+                return NULL;
+
+            PyObject *obj = PyUFunc_FromFuncAndData(
+                foo_loops, foo_data, foo_types, 1, 2, 1, PyUFunc_None, foo_name,
+                foo_docstring, 0);
+            if (!obj)
+            {
+                Py_DECREF(module);
+                return NULL;
+            }
+            if (PyModule_AddObject(module, foo_name, obj) < 0)
+            {
+                Py_DECREF(obj);
+                Py_DECREF(module);
+                return NULL;
+            }
+
+            return module;
+        }
+        """)
+    testdir.run(sys.executable, 'setup.py', 'build')
+    build_dir, = glob.glob(str(testdir.tmpdir / 'build/lib.*'))
+
+    result = testdir.inline_run(build_dir, '--doctest-plus', '--doctest-modules')
+    result.assertoutcome(passed=2, failed=0)