diff --git a/src/monitor.rs b/src/monitor.rs index 43e7159..092574a 100644 --- a/src/monitor.rs +++ b/src/monitor.rs @@ -26,12 +26,15 @@ pub struct Monitor { impl Monitor { /// Get The Primary Monitor - #[must_use] - pub fn primary() -> Self { + pub fn primary() -> Result> { let point = POINT { x: 0, y: 0 }; let monitor = unsafe { MonitorFromPoint(point, MONITOR_DEFAULTTOPRIMARY) }; - Self { monitor } + if monitor.is_invalid() { + return Err(Box::new(MonitorErrors::NotFound)); + } + + Ok(Self { monitor }) } /// Get The Monitor From It's Index diff --git a/windows-capture-python/rust-toolchain.toml b/windows-capture-python/rust-toolchain.toml new file mode 100644 index 0000000..fa60a1c --- /dev/null +++ b/windows-capture-python/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "stable" +targets = ["x86_64-pc-windows-msvc"] diff --git a/windows-capture-python/rustfmt.toml b/windows-capture-python/rustfmt.toml new file mode 100644 index 0000000..b34697f --- /dev/null +++ b/windows-capture-python/rustfmt.toml @@ -0,0 +1,12 @@ +wrap_comments = true +edition = "2021" +format_code_in_doc_comments = true +format_macro_matchers = true +format_strings = true +version = "Two" +use_try_shorthand = true +use_field_init_shorthand = true +unstable_features = true +normalize_doc_attributes = true +normalize_comments = true +imports_granularity = "Crate" diff --git a/windows-capture-python/src/lib.rs b/windows-capture-python/src/lib.rs index 13fc1a5..08924da 100644 --- a/windows-capture-python/src/lib.rs +++ b/windows-capture-python/src/lib.rs @@ -7,7 +7,7 @@ #![warn(clippy::cargo)] #![allow(clippy::redundant_pub_crate)] -use std::sync::Arc; +use std::{error::Error, sync::Arc}; use ::windows_capture::{ capture::WindowsCaptureHandler, @@ -15,6 +15,7 @@ use ::windows_capture::{ graphics_capture_api::InternalCaptureControl, monitor::Monitor, settings::{ColorFormat, WindowsCaptureSettings}, + window::Window, }; use log::{error, info}; use pyo3::{exceptions::PyException, prelude::*, types::PyList}; @@ -31,54 +32,113 @@ fn windows_capture(_py: Python, m: &PyModule) -> PyResult<()> { /// Internal Struct Used For Windows Capture #[pyclass] pub struct NativeWindowsCapture { - capture_cursor: bool, - draw_border: bool, on_frame_arrived_callback: Arc, on_closed: Arc, + capture_cursor: bool, + draw_border: bool, + monitor_index: Option, + window_name: Option, } #[pymethods] impl NativeWindowsCapture { /// Create A New Windows Capture Struct #[new] - #[must_use] pub fn new( - capture_cursor: bool, - draw_border: bool, on_frame_arrived_callback: PyObject, on_closed: PyObject, - ) -> Self { - Self { - capture_cursor, - draw_border, + capture_cursor: bool, + draw_border: bool, + monitor_index: Option, + window_name: Option, + ) -> PyResult { + if window_name.is_some() && monitor_index.is_some() { + return Err(PyException::new_err( + "You Can't Specify Both The Monitor Index And The Window Name", + )); + } + + if window_name.is_none() && monitor_index.is_none() { + return Err(PyException::new_err( + "You Should Specify Either The Monitor Index Or The Window Name", + )); + } + + Ok(Self { on_frame_arrived_callback: Arc::new(on_frame_arrived_callback), on_closed: Arc::new(on_closed), - } + capture_cursor, + draw_border, + monitor_index, + window_name, + }) } /// Start Capture pub fn start(&mut self) -> PyResult<()> { - let settings = match WindowsCaptureSettings::new( - Monitor::primary(), - Some(self.capture_cursor), - Some(self.draw_border), - ColorFormat::Bgra8, - ( - self.on_frame_arrived_callback.clone(), - self.on_closed.clone(), - ), - ) { - Ok(settings) => settings, - Err(e) => Err(PyException::new_err(format!( - "Failed To Create Windows Capture Settings -> {e}" - )))?, + let settings = if self.window_name.is_some() { + let window = match Window::from_contains_name(self.window_name.as_ref().unwrap()) { + Ok(window) => window, + Err(e) => { + return Err(PyException::new_err(format!( + "Failed To Find Window -> {e}" + ))); + } + }; + + match WindowsCaptureSettings::new( + window, + Some(self.capture_cursor), + Some(self.draw_border), + ColorFormat::Bgra8, + ( + self.on_frame_arrived_callback.clone(), + self.on_closed.clone(), + ), + ) { + Ok(settings) => settings, + Err(e) => { + return Err(PyException::new_err(format!( + "Failed To Create Windows Capture Settings -> {e}" + ))); + } + } + } else { + let monitor = match Monitor::from_index(self.monitor_index.unwrap()) { + Ok(monitor) => monitor, + Err(e) => { + return Err(PyException::new_err(format!( + "Failed To Get Monitor From Index -> {e}" + ))); + } + }; + + match WindowsCaptureSettings::new( + monitor, + Some(self.capture_cursor), + Some(self.draw_border), + ColorFormat::Bgra8, + ( + self.on_frame_arrived_callback.clone(), + self.on_closed.clone(), + ), + ) { + Ok(settings) => settings, + Err(e) => { + return Err(PyException::new_err(format!( + "Failed To Create Windows Capture Settings -> {e}" + ))); + } + } }; match InnerNativeWindowsCapture::start(settings) { Ok(_) => (), - Err(e) => Err(PyException::new_err(format!( - "Capture Session Threw An Exception -> {e}" - )))?, + Err(e) => { + return Err(PyException::new_err(format!( + "Capture Session Threw An Exception -> {e}" + ))); + } } Ok(()) @@ -93,14 +153,20 @@ struct InnerNativeWindowsCapture { impl WindowsCaptureHandler for InnerNativeWindowsCapture { type Flags = (Arc, Arc); - fn new((on_frame_arrived_callback, on_closed): Self::Flags) -> Self { - Self { + fn new( + (on_frame_arrived_callback, on_closed): Self::Flags, + ) -> Result> { + Ok(Self { on_frame_arrived_callback, on_closed, - } + }) } - fn on_frame_arrived(&mut self, mut frame: Frame, capture_control: InternalCaptureControl) { + fn on_frame_arrived( + &mut self, + mut frame: Frame, + capture_control: InternalCaptureControl, + ) -> Result<(), Box<(dyn Error + Send + Sync)>> { let width = frame.width(); let height = frame.height(); let buf = match frame.buffer() { @@ -110,48 +176,41 @@ impl WindowsCaptureHandler for InnerNativeWindowsCapture { "Failed To Get Frame Buffer -> {e} -> Gracefully Stopping The Capture Thread" ); capture_control.stop(); - return; + return Ok(()); } }; let buf = buf.as_raw_buffer(); - Python::with_gil(|py| { + Python::with_gil(|py| -> PyResult<()> { match py.check_signals() { Ok(_) => (), Err(_) => { info!("KeyboardInterrupt Detected -> Gracefully Stopping The Capture Thread"); capture_control.stop(); - return; + return Ok(()); } } let stop_list = PyList::new(py, [false]); - match self.on_frame_arrived_callback.call1( + self.on_frame_arrived_callback.call1( py, (buf.as_ptr() as isize, buf.len(), width, height, stop_list), - ) { - Ok(_) => (), - Err(e) => { - error!( - "on_frame_arrived Threw An Exception -> {e} -> Gracefully Stopping The \ - Capture Thread" - ); - capture_control.stop(); - return; - } - }; + )?; - if stop_list[0].is_true().unwrap_or(false) { + if stop_list[0].is_true()? { capture_control.stop(); } - }); + + Ok(()) + })?; + + Ok(()) } - fn on_closed(&mut self) { - Python::with_gil(|py| match self.on_closed.call0(py) { - Ok(_) => (), - Err(e) => error!("on_closed Threw An Exception -> {e}"), - }); + fn on_closed(&mut self) -> Result<(), Box<(dyn Error + Send + Sync)>> { + Python::with_gil(|py| self.on_closed.call0(py))?; + + Ok(()) } } diff --git a/windows-capture-python/windows_capture/__init__.py b/windows-capture-python/windows_capture/__init__.py index 6559d0e..1c8742d 100644 --- a/windows-capture-python/windows_capture/__init__.py +++ b/windows-capture-python/windows_capture/__init__.py @@ -5,6 +5,7 @@ import numpy import cv2 import types +from typing import Optional class Frame: @@ -19,13 +20,19 @@ class Frame: Raw Buffer Of The Frame width : str Width Of The Frame - age : int + height : int Height Of The Frame Methods ------- save_as_image(path: str): Saves The Frame As An Image To Specified Path + to_bgr() -> "Frame": + Converts The self.frame_buffer Pixel Type To Bgr Instead Of Bgra + crop( + start_width : int, start_height : int, end_width : int, end_height : int + ) -> "Frame": + Converts The self.frame_buffer Pixel Type To Bgr Instead Of Bgra """ def __init__(self, frame_buffer: numpy.ndarray, width: int, height: int) -> None: @@ -34,10 +41,28 @@ def __init__(self, frame_buffer: numpy.ndarray, width: int, height: int) -> None self.width = width self.height = height - def save_as_image(self, path: str): + def save_as_image(self, path: str) -> None: """Save The Frame As An Image To Specified Path""" cv2.imwrite(path, self.frame_buffer) + def convert_to_bgr(self) -> "Frame": + """Converts The self.frame_buffer Pixel Type To Bgr Instead Of Bgra""" + bgr_frame_buffer = self.frame_buffer[:, :, :3] + + return Frame(bgr_frame_buffer, self.width, self.height) + + def crop( + self, start_width: int, start_height: int, end_width: int, end_height: int + ) -> "Frame": + """Crops The Frame To The Specified Region""" + cropped_frame_buffer = self.frame_buffer[ + start_height:end_height, start_width:end_width, : + ] + + return Frame( + cropped_frame_buffer, end_width - start_width, end_height - start_height + ) + class CaptureControl: """ @@ -45,11 +70,6 @@ class CaptureControl: ... - Attributes - ---------- - _list : list - The First Index Is Used To Stop The Capture Thread - Methods ------- stop(): @@ -66,14 +86,80 @@ def stop(self) -> None: class WindowsCapture: - def __init__(self, capture_cursor: bool = True, draw_border: bool = False) -> None: - self.frame_handler = None - self.closed_handler = None + """ + Class To Capture The Screen + + ... + + Attributes + ---------- + frame_handler : Optional[types.FunctionType] + The on_frame_arrived Callback Function use @event to Override It Although It Can + Be Manually Changed + closed_handler : Optional[types.FunctionType] + The on_closed Callback Function use @event to Override It Although It Can Be + Manually Changed + + Methods + ------- + start(): + Starts The Capture Thread + on_frame_arrived( + buf : ctypes.POINTER, + buf_len : int, + width : int, + height : int, + stop_list : list, + ): + This Method Is Called Before The on_frame_arrived Callback Function NEVER + Modify This Method Only Modify The Callback AKA frame_handler + on_closed(): + This Method Is Called Before The on_closed Callback Function To + Prepare Data NEVER Modify This Method + Only Modify The Callback AKA closed_handler + event(handler: types.FunctionType): + Overrides The Callback Function + """ + + def __init__( + self, + capture_cursor: bool = True, + draw_border: bool = False, + monitor_index: int = 0, + window_name: Optional[str] = None, + ) -> None: + """ + Constructs All The Necessary Attributes For The WindowsCapture Object + + ... + + Parameters + ---------- + capture_cursor : bool + Whether To Capture The Cursor + draw_border : bool + Whether To draw The border + monitor_index : int + Index Of The Monitor To Capture + window_name : str + Name Of The Window To Capture + """ + if window_name is not None: + monitor_index = None + + self.frame_handler: Optional[types.FunctionType] = None + self.closed_handler: Optional[types.FunctionType] = None self.capture = NativeWindowsCapture( - capture_cursor, draw_border, self.on_frame_arrived, self.on_closed + self.on_frame_arrived, + self.on_closed, + capture_cursor, + draw_border, + monitor_index, + window_name, ) def start(self) -> None: + """Starts The Capture Thread""" if self.frame_handler is None: raise Exception("on_frame_arrived Event Handler Is Not Set") elif self.closed_handler is None: @@ -89,25 +175,27 @@ def on_frame_arrived( height: int, stop_list: list, ) -> None: + """This Method Is Called Before The on_frame_arrived Callback Function To + Prepare Data""" if self.frame_handler: internal_capture_control = CaptureControl(stop_list) row_pitch = buf_len / height if row_pitch == width * 4: - num_array = numpy.ctypeslib.as_array( + ndarray = numpy.ctypeslib.as_array( ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8)), shape=(height, width, 4), ) - frame = Frame(num_array, width, height) + frame = Frame(ndarray, width, height) self.frame_handler(frame, internal_capture_control) else: - num_array = numpy.ctypeslib.as_array( + ndarray = numpy.ctypeslib.as_array( ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8)), shape=(height, row_pitch), )[:, : width * 4].reshape(height, width, 4) - frame = Frame(num_array, width, height) + frame = Frame(ndarray, width, height) self.frame_handler(frame, internal_capture_control) self.frame_handler( @@ -119,12 +207,14 @@ def on_frame_arrived( raise Exception("on_frame_arrived Event Handler Is Not Set") def on_closed(self) -> None: + """This Method Is Called Before The on_closed Callback Function""" if self.closed_handler: self.closed_handler() else: raise Exception("on_closed Event Handler Is Not Set") def event(self, handler: types.FunctionType) -> None: + """Overrides The Callback Function""" if handler.__name__ == "on_frame_arrived": self.frame_handler = handler elif handler.__name__ == "on_closed":