diff --git a/python/pylibraft/pylibraft/common/__init__.py b/python/pylibraft/pylibraft/common/__init__.py index 33c2986487..9c0f631b86 100644 --- a/python/pylibraft/pylibraft/common/__init__.py +++ b/python/pylibraft/pylibraft/common/__init__.py @@ -17,3 +17,5 @@ from .cuda import Stream from .device_ndarray import device_ndarray from .handle import Handle + +__all__ = ["Handle", "Stream"] diff --git a/python/pylibraft/pylibraft/common/handle.pyx b/python/pylibraft/pylibraft/common/handle.pyx index f662f5bd70..13fc7fc98e 100644 --- a/python/pylibraft/pylibraft/common/handle.pyx +++ b/python/pylibraft/pylibraft/common/handle.pyx @@ -37,7 +37,7 @@ cdef class Handle: Examples -------- - >>> from raft.common import Stream, Handle + >>> from pylibraft.common import Stream, Handle >>> stream = Stream() >>> handle = Handle(stream) >>> diff --git a/python/pylibraft/pylibraft/test/test_doctests.py b/python/pylibraft/pylibraft/test/test_doctests.py index 29260c7128..3276ca115f 100644 --- a/python/pylibraft/pylibraft/test/test_doctests.py +++ b/python/pylibraft/pylibraft/test/test_doctests.py @@ -92,6 +92,7 @@ def _find_doctests_in_obj(obj, finder=None, criteria=None): # __all__) we are explicitly adding all the submodules we want to run # doctests for here DOC_STRINGS = list(_find_doctests_in_obj(pylibraft.cluster)) +DOC_STRINGS.extend(_find_doctests_in_obj(pylibraft.common)) DOC_STRINGS.extend(_find_doctests_in_obj(pylibraft.distance)) DOC_STRINGS.extend(_find_doctests_in_obj(pylibraft.neighbors)) DOC_STRINGS.extend(_find_doctests_in_obj(pylibraft.neighbors.ivf_pq))