Update bitmap functions to accept hex value rather than tuple

autopilot-rs · Feb 26, 2020 · 0b60f22 · 0b60f22
1 parent b1e0e03
commit 0b60f22
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 17 deletions.
diff --git a/src/bitmap.rs b/src/bitmap.rs
@@ -9,7 +9,7 @@ use autopilot::geometry::{Point, Rect, Size};
 use image;
 use image::Pixel;
 use image::{ImageOutputFormat, ImageResult, Rgba};
-use internal::{rgb_to_hex, FromImageError};
+use internal::{rgb_to_hex, hex_to_rgb, FromImageError};
 use pyo3::basic::CompareOp;
 use pyo3::prelude::*;
 use pyo3::types::PyType;
@@ -191,16 +191,17 @@ impl Bitmap {
     /// exact match and 1 matches anything.
     fn find_color(
         &self,
-        color: (u8, u8, u8),
+        color: u32,
         tolerance: Option<f64>,
         rect: Option<((f64, f64), (f64, f64))>,
         start_point: Option<(f64, f64)>,
     ) -> PyResult<Option<(f64, f64)>> {
-        let color = Rgba([color.0, color.1, color.2, 255]);
+        let split_color = hex_to_rgb(color);
+        let rgb = Rgba([split_color.0, split_color.1, split_color.2, 255]);
         let rect: Option<Rect> =
             rect.map(|r| Rect::new(Point::new((r.0).0, (r.0).1), Size::new((r.1).0, (r.1).1)));
         let start_point: Option<Point> = start_point.map(|p| Point::new(p.0, p.1));
-        if let Some(point) = self.bitmap.find_color(color, tolerance, rect, start_point) {
+        if let Some(point) = self.bitmap.find_color(rgb, tolerance, rect, start_point) {
             Ok(Some((point.x, point.y)))
         } else {
             Ok(None)
@@ -213,18 +214,19 @@ impl Bitmap {
     /// `rect` is used.
     fn find_every_color(
         &self,
-        color: (u8, u8, u8),
+        color: u32,
         tolerance: Option<f64>,
         rect: Option<((f64, f64), (f64, f64))>,
         start_point: Option<(f64, f64)>,
     ) -> PyResult<Vec<(f64, f64)>> {
-        let color = Rgba([color.0, color.1, color.2, 255]);
+        let split_color = hex_to_rgb(color);
+        let rgb = Rgba([split_color.0, split_color.1, split_color.2, 255]);
         let rect: Option<Rect> =
             rect.map(|r| Rect::new(Point::new((r.0).0, (r.0).1), Size::new((r.1).0, (r.1).1)));
         let start_point: Option<Point> = start_point.map(|p| Point::new(p.0, p.1));
         let points = self
             .bitmap
-            .find_every_color(color, tolerance, rect, start_point)
+            .find_every_color(rgb, tolerance, rect, start_point)
             .iter()
             .map(|p| (p.x, p.y))
             .collect();
@@ -236,18 +238,19 @@ impl Bitmap {
     /// `len(find_every_color(color, tolerance, rect, start_point))`
     fn count_of_color(
         &self,
-        color: (u8, u8, u8),
+        color: u32,
         tolerance: Option<f64>,
         rect: Option<((f64, f64), (f64, f64))>,
         start_point: Option<(f64, f64)>,
     ) -> PyResult<u64> {
-        let color = Rgba([color.0, color.1, color.2, 255]);
+        let split_color = hex_to_rgb(color);
+        let rgb = Rgba([split_color.0, split_color.1, split_color.2, 255]);
         let rect: Option<Rect> =
             rect.map(|r| Rect::new(Point::new((r.0).0, (r.0).1), Size::new((r.1).0, (r.1).1)));
         let start_point: Option<Point> = start_point.map(|p| Point::new(p.0, p.1));
         let count = self
             .bitmap
-            .count_of_color(color, tolerance, rect, start_point);
+            .count_of_color(rgb, tolerance, rect, start_point);
         Ok(count)
     }
 

diff --git a/src/color.rs b/src/color.rs
@@ -5,7 +5,7 @@
 // https://opensource.org/licenses/MIT>, at your option. This file may not be
 // copied, modified, or distributed except according to those terms.
 
-use internal::rgb_to_hex;
+use internal::{rgb_to_hex, hex_to_rgb};
 use pyo3::prelude::*;
 
 /// Returns hexadecimal value of given RGB tuple. `r`, `g`, and `b` must be in
@@ -18,11 +18,8 @@ fn py_rgb_to_hex(red: u8, green: u8, blue: u8) -> PyResult<u32> {
 /// Returns a tuple `(r, g, b)` of the RGB integer values equivalent to the
 /// given RGB hexadecimal value. `r`, `g`, and `b` are in the range 0 - 255.
 #[pyfunction]
-fn hex_to_rgb(hex: u32) -> PyResult<(u8, u8, u8)> {
-    let red: u8 = ((hex >> 16) & 0xff) as u8;
-    let green: u8 = ((hex >> 8) & 0xff) as u8;
-    let blue: u8 = (hex & 0xff) as u8;
-    Ok((red, green, blue))
+fn py_hex_to_rgb(hex: u32) -> PyResult<(u8, u8, u8)> {
+    Ok(hex_to_rgb(hex))
 }
 
 /// This module provides functions for converting between the hexadecimal
@@ -31,6 +28,6 @@ fn hex_to_rgb(hex: u32) -> PyResult<(u8, u8, u8)> {
 #[pymodule(color)]
 fn init(py: Python, m: &PyModule) -> PyResult<()> {
     m.add("rgb_to_hex", wrap_pyfunction!(py_rgb_to_hex)(py))?;
-    m.add_wrapped(wrap_pyfunction!(hex_to_rgb))?;
+    m.add("hex_to_rgb", wrap_pyfunction!(py_hex_to_rgb)(py))?;
     Ok(())
 }
diff --git a/src/internal.rs b/src/internal.rs
@@ -9,10 +9,18 @@ use image::ImageError;
 use pyo3::prelude::*;
 
 pub struct FromImageError(ImageError);
+
 pub fn rgb_to_hex(red: u8, green: u8, blue: u8) -> u32 {
     ((red as u32) << 16) | ((green as u32) << 8) | blue as u32
 }
 
+pub fn hex_to_rgb(hex: u32) -> (u8, u8, u8) {
+    let red: u8 = ((hex >> 16) & 0xff) as u8;
+    let green: u8 = ((hex >> 8) & 0xff) as u8;
+    let blue: u8 = (hex & 0xff) as u8;
+    (red, green, blue)
+}
+
 impl From<ImageError> for FromImageError {
     fn from(err: ImageError) -> FromImageError {
         FromImageError { 0: err }