improvement: add numpy support to range_slider, update examples

marimo/_plugins/ui/_impl/input.py

@@ -167,6 +167,20 @@ class slider(UIElement[Numeric, Numeric]):
     slider = mo.ui.slider.from_series(df["column_name"])
     ```
 
+    Or using numpy arrays:
+
+    ```python
+    import numpy as np
+
+    # linear steps
+    steps = np.array([1, 2, 3, 4, 5])
+    slider = mo.ui.slider(steps=steps)
+    # log steps
+    log_slider = mo.ui.slider(steps=np.logspace(0, 3, 4))
+    # power steps
+    power_slider = mo.ui.slider(steps=np.power([1, 2, 3], 2))
+    ```
+
     **Attributes.**
 
     - `value`: the current numeric value of the slider
@@ -235,11 +249,7 @@ def __init__(
         if steps is not None:
             # Cast to a list in case user passes a numpy array
             if not isinstance(steps, list):
-                if DependencyManager.numpy.has():
-                    import numpy as np
-
-                    if isinstance(steps, np.ndarray):
-                        steps = steps.tolist()
+                steps = _convert_numpy_array(steps)
             self._dtype = _infer_dtype(steps)
             self._mapping = dict(enumerate(steps))
             try:
@@ -358,6 +368,20 @@ class range_slider(UIElement[List[Numeric], Sequence[Numeric]]):
     range_slider = mo.ui.range_slider.from_series(df["column_name"])
     ```
 
+    Or using numpy arrays:
+
+    ```python
+    import numpy as np
+
+    steps = np.array([1, 2, 3, 4, 5])
+    # linear steps
+    range_slider = mo.ui.range_slider(steps=steps)
+    # log steps
+    log_range_slider = mo.ui.range_slider(steps=np.logspace(0, 3, 4))
+    # power steps
+    power_range_slider = mo.ui.range_slider(steps=np.power([1, 2, 3], 2))
+    ```
+
     **Attributes.**
 
     - `value`: the current range value of the slider
@@ -423,6 +447,9 @@ def __init__(
             )
 
         if steps is not None:
+            # Cast to a list in case user passes a numpy array
+            if not isinstance(steps, list):
+                steps = _convert_numpy_array(steps)
             self._dtype = _infer_dtype(steps)
             self._mapping = dict(enumerate(steps))
 
@@ -1619,3 +1646,13 @@ def _convert_value(self, value: Optional[JSONTypeBound]) -> Optional[T]:
             return None
         self.element._update(value)
         return self.element.value
+
+
+def _convert_numpy_array(steps: Any) -> list[Numeric]:
+    """Convert numpy array to list if needed."""
+    if DependencyManager.numpy.has():
+        import numpy as np
+
+        if isinstance(steps, np.ndarray):
+            return steps.tolist()  # type: ignore[no-any-return]
+    return steps  # type: ignore[no-any-return]

tests/_plugins/ui/_impl/test_input.py

@@ -281,15 +281,6 @@ def test_range_slider_from_dataframe() -> None:
     assert slider._args.label == "Custom label"
 
 
-@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not installed")
-def test_range_slider_from_np_array() -> None:
-    import numpy as np
-
-    slider = ui.slider(steps=np.array([1, 2, 3]))
-    assert slider.start == 1
-    assert slider.stop == 3
-
-
 def test_text() -> None:
     assert ui.text().value == ""
     assert ui.text(value="hello world").value == "hello world"
@@ -470,3 +461,58 @@ def test_form_in_dictionary_allowed() -> None:
 
 
 # TODO(akshayka): test file
+
+
+@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not installed")
+def test_numpy_steps() -> None:
+    import numpy as np
+
+    steps = np.array([1, 2, 3, 4, 5])
+    slider = ui.slider(steps=steps)
+    assert slider.steps == [1, 2, 3, 4, 5]
+    assert slider.start == 1
+    assert slider.stop == 5
+    assert slider.step is None
+
+    range_slider = ui.range_slider(steps=steps)
+    assert range_slider.steps == [1, 2, 3, 4, 5]
+    assert range_slider.start == 1
+    assert range_slider.stop == 5
+    assert range_slider.step is None
+
+
+@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not installed")
+def test_log_scale() -> None:
+    import numpy as np
+
+    steps = np.logspace(0, 3, 4)
+    slider = ui.slider(steps=steps)
+
+    assert slider.steps == [1, 10, 100, 1000]
+    assert slider.start == 1
+    assert slider.stop == 1000
+    assert slider.step is None
+
+    range_slider = ui.range_slider(steps=steps)
+    assert range_slider.steps == [1, 10, 100, 1000]
+    assert range_slider.start == 1
+    assert range_slider.stop == 1000
+    assert range_slider.step is None
+
+
+@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not installed")
+def test_power_scale() -> None:
+    import numpy as np
+
+    steps = np.power([1, 2, 3], 2)
+    slider = ui.slider(steps=steps)
+    assert slider.steps == [1, 4, 9]
+    assert slider.start == 1
+    assert slider.stop == 9
+    assert slider.step is None
+
+    range_slider = ui.range_slider(steps=steps)
+    assert range_slider.steps == [1, 4, 9]
+    assert range_slider.start == 1
+    assert range_slider.stop == 9
+    assert range_slider.step is None