feat: support color picker functionality

chore: format code

refactor: improve quality

feat: implement gnomes PickColor method

refactor: minor code extraction

misc: fix reviews

fixes

niri-ipc/src/lib.rs

@@ -65,6 +65,8 @@ pub enum Request {
     FocusedWindow,
     /// Request picking a window and get its information.
     PickWindow,
+    /// Request picking a color from the screen.
+    PickColor,
     /// Perform an action.
     Action(Action),
     /// Change output configuration temporarily.
@@ -133,10 +135,20 @@ pub enum Response {
     FocusedWindow(Option<Window>),
     /// Information about the picked window.
     PickedWindow(Option<Window>),
+    /// Information about the picked color.
+    PickedColor(Option<PickedColor>),
     /// Output configuration change result.
     OutputConfigChanged(OutputConfigChanged),
 }
 
+/// Color picked from the screen.
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
+pub struct PickedColor {
+    /// Color values as red, green, blue, each ranging from 0.0 to 1.0.
+    pub rgb: [f64; 3],
+}
+
 /// Actions that niri can perform.
 // Variants in this enum should match the spelling of the ones in niri-config. Most, but not all,
 // variants from niri-config should be present here.

src/cli.rs

@@ -77,6 +77,8 @@ pub enum Msg {
     FocusedWindow,
     /// Pick a window with the mouse and print information about it.
     PickWindow,
+    /// Pick a color from the screen with the mouse.
+    PickColor,
     /// Perform an action.
     Action {
         #[command(subcommand)]

src/dbus/gnome_shell_screenshot.rs

@@ -1,7 +1,10 @@
+use std::collections::HashMap;
 use std::path::PathBuf;
 
+use niri_ipc::PickedColor;
 use zbus::fdo::{self, RequestNameFlags};
 use zbus::interface;
+use zbus::zvariant::OwnedValue;
 
 use super::Start;
 
@@ -12,6 +15,7 @@ pub struct Screenshot {
 
 pub enum ScreenshotToNiri {
     TakeScreenshot { include_cursor: bool },
+    PickColor(async_channel::Sender<Option<PickedColor>>),
 }
 
 pub enum NiriToScreenshot {
@@ -47,6 +51,34 @@ impl Screenshot {
 
         Ok((true, filename))
     }
+
+    async fn pick_color(&self) -> fdo::Result<HashMap<String, OwnedValue>> {
+        let (tx, rx) = async_channel::bounded(1);
+        if let Err(err) = self.to_niri.send(ScreenshotToNiri::PickColor(tx)) {
+            warn!("error sending pick color message to niri: {err:?}");
+            return Err(fdo::Error::Failed("internal error".to_owned()));
+        }
+
+        let color = match rx.recv().await {
+            Ok(Some(color)) => color,
+            Ok(None) => {
+                return Err(fdo::Error::Failed("no color picked".to_owned()));
+            }
+            Err(err) => {
+                warn!("error receiving message from niri: {err:?}");
+                return Err(fdo::Error::Failed("internal error".to_owned()));
+            }
+        };
+
+        let mut result = HashMap::new();
+        let rgb_slice: &[f64] = &color.rgb;
+        result.insert(
+            "color".to_string(),
+            zbus::zvariant::Value::from(rgb_slice).try_into().unwrap(),
+        );
+
+        Ok(result)
+    }
 }
 
 impl Screenshot {

src/input/mod.rs

@@ -48,6 +48,7 @@ use crate::utils::{center, get_monotonic_time, ResizeEdge};
 
 pub mod backend_ext;
 pub mod move_grab;
+pub mod pick_color_grab;
 pub mod pick_window_grab;
 pub mod resize_grab;
 pub mod scroll_tracker;
@@ -377,7 +378,10 @@ impl State {
                     }
                 }
 
-                if pressed && raw == Some(Keysym::Escape) && this.niri.pick_window.is_some() {
+                if pressed
+                    && raw == Some(Keysym::Escape)
+                    && (this.niri.pick_window.is_some() || this.niri.pick_color.is_some())
+                {
                     // We window picking state so the pick window grab must be active.
                     // Unsetting it cancels window picking.
                     this.niri

src/input/pick_color_grab.rs

@@ -0,0 +1,223 @@
+use niri_ipc::PickedColor;
+use smithay::backend::allocator::Fourcc;
+use smithay::backend::input::ButtonState;
+use smithay::backend::renderer::element::utils::{Relocate, RelocateRenderElement};
+use smithay::input::pointer::{
+    AxisFrame, ButtonEvent, CursorImageStatus, GestureHoldBeginEvent, GestureHoldEndEvent,
+    GesturePinchBeginEvent, GesturePinchEndEvent, GesturePinchUpdateEvent, GestureSwipeBeginEvent,
+    GestureSwipeEndEvent, GestureSwipeUpdateEvent, GrabStartData as PointerGrabStartData,
+    MotionEvent, PointerGrab, PointerInnerHandle, RelativeMotionEvent,
+};
+use smithay::input::SeatHandler;
+use smithay::utils::{Logical, Physical, Point, Scale, Size, Transform};
+
+use crate::niri::State;
+use crate::render_helpers::render_to_vec;
+
+pub struct PickColorGrab {
+    start_data: PointerGrabStartData<State>,
+}
+
+impl PickColorGrab {
+    pub fn new(start_data: PointerGrabStartData<State>) -> Self {
+        Self { start_data }
+    }
+
+    fn on_ungrab(&mut self, state: &mut State) {
+        if let Some(tx) = state.niri.pick_color.take() {
+            let _ = tx.send_blocking(None);
+        }
+        state
+            .niri
+            .cursor_manager
+            .set_cursor_image(CursorImageStatus::default_named());
+        state.niri.queue_redraw_all();
+    }
+
+    fn pick_color_at_point(location: Point<f64, Logical>, data: &mut State) -> Option<PickedColor> {
+        let (output, pos_within_output) = data.niri.output_under(location)?;
+        let output = output.clone();
+
+        data.backend
+            .with_primary_renderer(|renderer| {
+                data.niri.update_render_elements(Some(&output));
+
+                let scale = Scale::from(output.current_scale().fractional_scale());
+                // FIXME: perhaps replace floor with round once we figure out the pointer behavior
+                // at the bottom/right edges of the monitors.
+                let pos = pos_within_output.to_physical_precise_floor(scale);
+                let size = Size::<i32, Physical>::from((1, 1));
+
+                let elements = data.niri.render(
+                    renderer,
+                    &output,
+                    false,
+                    crate::render_helpers::RenderTarget::ScreenCapture,
+                );
+
+                let pixels = match render_to_vec(
+                    renderer,
+                    size,
+                    scale,
+                    Transform::Normal,
+                    Fourcc::Abgr8888,
+                    elements.iter().rev().map(|elem| {
+                        let offset = pos.upscale(-1);
+                        RelocateRenderElement::from_element(elem, offset, Relocate::Relative)
+                    }),
+                ) {
+                    Ok(pixels) => pixels,
+                    Err(_) => return None,
+                };
+
+                if pixels.len() == 4 {
+                    let rgb = [
+                        f64::from(pixels[0]) / 255.0,
+                        f64::from(pixels[1]) / 255.0,
+                        f64::from(pixels[2]) / 255.0,
+                    ];
+                    Some(PickedColor { rgb })
+                } else {
+                    error!(
+                        "unexpected pixel data length: {} (expected 4)",
+                        pixels.len()
+                    );
+                    None
+                }
+            })
+            .flatten()
+    }
+}
+
+impl PointerGrab<State> for PickColorGrab {
+    fn motion(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        _focus: Option<(<State as SeatHandler>::PointerFocus, Point<f64, Logical>)>,
+        event: &MotionEvent,
+    ) {
+        handle.motion(data, None, event);
+    }
+
+    fn relative_motion(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        _focus: Option<(<State as SeatHandler>::PointerFocus, Point<f64, Logical>)>,
+        event: &RelativeMotionEvent,
+    ) {
+        handle.relative_motion(data, None, event);
+    }
+
+    fn button(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &ButtonEvent,
+    ) {
+        if event.state != ButtonState::Pressed {
+            return;
+        }
+
+        if let Some(tx) = data.niri.pick_color.take() {
+            let color = Self::pick_color_at_point(handle.current_location(), data);
+            let _ = tx.send_blocking(color);
+        }
+
+        handle.unset_grab(self, data, event.serial, event.time, true);
+    }
+
+    fn axis(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        details: AxisFrame,
+    ) {
+        handle.axis(data, details);
+    }
+
+    fn frame(&mut self, data: &mut State, handle: &mut PointerInnerHandle<'_, State>) {
+        handle.frame(data);
+    }
+
+    fn gesture_swipe_begin(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GestureSwipeBeginEvent,
+    ) {
+        handle.gesture_swipe_begin(data, event);
+    }
+
+    fn gesture_swipe_update(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GestureSwipeUpdateEvent,
+    ) {
+        handle.gesture_swipe_update(data, event);
+    }
+
+    fn gesture_swipe_end(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GestureSwipeEndEvent,
+    ) {
+        handle.gesture_swipe_end(data, event);
+    }
+
+    fn gesture_pinch_begin(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GesturePinchBeginEvent,
+    ) {
+        handle.gesture_pinch_begin(data, event);
+    }
+
+    fn gesture_pinch_update(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GesturePinchUpdateEvent,
+    ) {
+        handle.gesture_pinch_update(data, event);
+    }
+
+    fn gesture_pinch_end(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GesturePinchEndEvent,
+    ) {
+        handle.gesture_pinch_end(data, event);
+    }
+
+    fn gesture_hold_begin(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GestureHoldBeginEvent,
+    ) {
+        handle.gesture_hold_begin(data, event);
+    }
+
+    fn gesture_hold_end(
+        &mut self,
+        data: &mut State,
+        handle: &mut PointerInnerHandle<'_, State>,
+        event: &GestureHoldEndEvent,
+    ) {
+        handle.gesture_hold_end(data, event);
+    }
+
+    fn start_data(&self) -> &PointerGrabStartData<State> {
+        &self.start_data
+    }
+
+    fn unset(&mut self, data: &mut State) {
+        self.on_ungrab(data);
+    }
+}

src/ipc/client.rs

@@ -20,6 +20,7 @@ pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
         Msg::FocusedWindow => Request::FocusedWindow,
         Msg::FocusedOutput => Request::FocusedOutput,
         Msg::PickWindow => Request::PickWindow,
+        Msg::PickColor => Request::PickColor,
         Msg::Action { action } => Request::Action(action.clone()),
         Msg::Output { output, action } => Request::Output {
             output: output.clone(),
@@ -270,6 +271,26 @@ pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
                 println!("No window selected.");
             }
         }
+        Msg::PickColor => {
+            let Response::PickedColor(color) = response else {
+                bail!("unexpected response: expected PickedColor, got {response:?}");
+            };
+
+            if json {
+                let color = serde_json::to_string(&color).context("error formatting response")?;
+                println!("{color}");
+                return Ok(());
+            }
+
+            if let Some(color) = color {
+                let [r, g, b] = color.rgb.map(|v| (v.clamp(0., 1.) * 255.).round() as u8);
+
+                println!("Picked color: rgb({r}, {g}, {b})",);
+                println!("Hex: #{:02x}{:02x}{:02x}", r, g, b);
+            } else {
+                println!("No color was picked.");
+            }
+        }
         Msg::Action { .. } => {
             let Response::Handled = response else {
                 bail!("unexpected response: expected Handled, got {response:?}");

src/ipc/server.rs

@@ -356,6 +356,15 @@ async fn process(ctx: &ClientCtx, request: Request) -> Reply {
             });
             Response::PickedWindow(window)
         }
+        Request::PickColor => {
+            let (tx, rx) = async_channel::bounded(1);
+            ctx.event_loop.insert_idle(move |state| {
+                state.handle_pick_color(tx);
+            });
+            let result = rx.recv().await;
+            let color = result.map_err(|_| String::from("error getting picked color"))?;
+            Response::PickedColor(color)
+        }
         Request::Action(action) => {
             let (tx, rx) = async_channel::bounded(1);
 

src/niri.rs

@@ -45,7 +45,10 @@ use smithay::desktop::{
     PopupUngrabStrategy, Space, Window, WindowSurfaceType,
 };
 use smithay::input::keyboard::Layout as KeyboardLayout;
-use smithay::input::pointer::{CursorIcon, CursorImageStatus, CursorImageSurfaceData, MotionEvent};
+use smithay::input::pointer::{
+    CursorIcon, CursorImageStatus, CursorImageSurfaceData, Focus,
+    GrabStartData as PointerGrabStartData, MotionEvent,
+};
 use smithay::input::{Seat, SeatState};
 use smithay::output::{self, Output, OutputModeSource, PhysicalProperties, Subpixel, WeakOutput};
 use smithay::reexports::calloop::generic::Generic;
@@ -118,6 +121,7 @@ use crate::dbus::gnome_shell_screenshot::{NiriToScreenshot, ScreenshotToNiri};
 use crate::dbus::mutter_screen_cast::{self, ScreenCastToNiri};
 use crate::frame_clock::FrameClock;
 use crate::handlers::{configure_lock_surface, XDG_ACTIVATION_TOKEN_TIMEOUT};
+use crate::input::pick_color_grab::PickColorGrab;
 use crate::input::scroll_tracker::ScrollTracker;
 use crate::input::{
     apply_libinput_settings, mods_with_finger_scroll_binds, mods_with_mouse_binds,
@@ -358,6 +362,7 @@ pub struct Niri {
     pub exit_confirm_dialog: Option<ExitConfirmDialog>,
 
     pub pick_window: Option<async_channel::Sender<Option<MappedId>>>,
+    pub pick_color: Option<async_channel::Sender<Option<niri_ipc::PickedColor>>>,
 
     pub debug_draw_opaque_regions: bool,
     pub debug_draw_damage: bool,
@@ -1612,6 +1617,22 @@ impl State {
         self.niri.queue_redraw_all();
     }
 
+    pub fn handle_pick_color(&mut self, tx: async_channel::Sender<Option<niri_ipc::PickedColor>>) {
+        let pointer = self.niri.seat.get_pointer().unwrap();
+        let start_data = PointerGrabStartData {
+            focus: None,
+            button: 0,
+            location: pointer.current_location(),
+        };
+        let grab = PickColorGrab::new(start_data);
+        pointer.set_grab(self, grab, SERIAL_COUNTER.next_serial(), Focus::Clear);
+        self.niri.pick_color = Some(tx);
+        self.niri
+            .cursor_manager
+            .set_cursor_image(CursorImageStatus::Named(CursorIcon::Crosshair));
+        self.niri.queue_redraw_all();
+    }
+
     #[cfg(feature = "xdp-gnome-screencast")]
     pub fn on_pw_msg(&mut self, msg: PwToNiri) {
         match msg {
@@ -1903,7 +1924,22 @@ impl State {
         to_screenshot: &async_channel::Sender<NiriToScreenshot>,
         msg: ScreenshotToNiri,
     ) {
-        let ScreenshotToNiri::TakeScreenshot { include_cursor } = msg;
+        match msg {
+            ScreenshotToNiri::TakeScreenshot { include_cursor } => {
+                self.handle_take_screenshot(to_screenshot, include_cursor);
+            }
+            ScreenshotToNiri::PickColor(tx) => {
+                self.handle_pick_color(tx);
+            }
+        }
+    }
+
+    #[cfg(feature = "dbus")]
+    fn handle_take_screenshot(
+        &mut self,
+        to_screenshot: &async_channel::Sender<NiriToScreenshot>,
+        include_cursor: bool,
+    ) {
         let _span = tracy_client::span!("TakeScreenshot");
 
         let rv = self.backend.with_primary_renderer(|renderer| {
@@ -2351,6 +2387,7 @@ impl Niri {
             exit_confirm_dialog,
 
             pick_window: None,
+            pick_color: None,
 
             debug_draw_opaque_regions: false,
             debug_draw_damage: false,