Add lazy loading to skimage.draw submodule (#6971)

scikit-image · May 31, 2023 · 0c17a8f · 0c17a8f
1 parent 907d272
commit 0c17a8f
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 30 deletions.
diff --git a/skimage/draw/__init__.py b/skimage/draw/__init__.py
@@ -1,29 +1,3 @@
-from .draw import (ellipse, set_color, polygon_perimeter,
-                   line, line_aa, polygon, ellipse_perimeter,
-                   circle_perimeter, circle_perimeter_aa,
-                   disk, bezier_curve, rectangle, rectangle_perimeter)
-from .draw3d import ellipsoid, ellipsoid_stats
-from ._draw import _bezier_segment
-from ._random_shapes import random_shapes
-from ._polygon2mask import polygon2mask
+import lazy_loader as lazy
 
-from .draw_nd import line_nd
-
-__all__ = ['line',
-           'line_aa',
-           'line_nd',
-           'bezier_curve',
-           'polygon',
-           'polygon_perimeter',
-           'ellipse',
-           'ellipse_perimeter',
-           'ellipsoid',
-           'ellipsoid_stats',
-           'circle_perimeter',
-           'circle_perimeter_aa',
-           'disk',
-           'set_color',
-           'random_shapes',
-           'rectangle',
-           'rectangle_perimeter',
-           'polygon2mask']
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)
diff --git a/skimage/draw/__init__.pyi b/skimage/draw/__init__.pyi
@@ -0,0 +1,44 @@
+# Explicitly setting `__all__` is necessary for type inference engines
+# to know which symbols are exported. See
+# https://peps.python.org/pep-0484/#stub-files
+
+__all__ = [
+    'line',
+    'line_aa',
+    'line_nd',
+    'bezier_curve',
+    'polygon',
+    'polygon_perimeter',
+    'ellipse',
+    'ellipse_perimeter',
+    'ellipsoid',
+    'ellipsoid_stats',
+    'circle_perimeter',
+    'circle_perimeter_aa',
+    'disk',
+    'set_color',
+    'random_shapes',
+    'rectangle',
+    'rectangle_perimeter',
+    'polygon2mask'
+]
+
+from .draw3d import ellipsoid, ellipsoid_stats
+from ._draw import _bezier_segment
+from ._random_shapes import random_shapes
+from ._polygon2mask import polygon2mask
+from .draw_nd import line_nd
+from .draw import (
+    ellipse,
+    set_color,
+    polygon_perimeter,
+    line,
+    line_aa,
+    polygon,
+    ellipse_perimeter,
+    circle_perimeter,
+    circle_perimeter_aa,
+    disk, bezier_curve,
+    rectangle,
+    rectangle_perimeter
+)
diff --git a/skimage/draw/_random_shapes.py b/skimage/draw/_random_shapes.py
@@ -2,8 +2,8 @@
 
 import numpy as np
 
-from . import (polygon as draw_polygon, disk as draw_disk,
-               ellipse as draw_ellipse)
+from .draw import (polygon as draw_polygon, disk as draw_disk,
+                   ellipse as draw_ellipse)
 from .._shared.utils import deprecate_kwarg, warn
 
 

diff --git a/skimage/draw/meson.build b/skimage/draw/meson.build
@@ -8,6 +8,7 @@ py3.extension_module('_draw',
 
 python_sources = [
   '__init__.py',
+  '__init__.pyi',
   '_polygon2mask.py',
   '_random_shapes.py',
   'draw3d.py',