diff --git a/src/capture.rs b/src/capture.rs index e1130ff..bba4504 100644 --- a/src/capture.rs +++ b/src/capture.rs @@ -10,6 +10,7 @@ use std::{ }; use log::{debug, info, trace, warn}; +use parking_lot::Mutex; use windows::{ Foundation::AsyncActionCompletedHandler, Win32::{ @@ -39,24 +40,29 @@ use crate::{ pub enum CaptureControlError { #[error("Failed To Join Thread")] FailedToJoinThread, + #[error("Thread Handle Is Taken Out Of Struct")] + ThreadHandleIsTaken, } /// Struct Used To Control Capture Thread -pub struct CaptureControl { +pub struct CaptureControl { thread_handle: Option>>>, halt_handle: Arc, + callback: Arc>, } -impl CaptureControl { +impl CaptureControl { /// Create A New Capture Control Struct #[must_use] pub fn new( thread_handle: JoinHandle>>, halt_handle: Arc, + callback: Arc>, ) -> Self { Self { thread_handle: Some(thread_handle), halt_handle, + callback, } } @@ -68,6 +74,24 @@ impl CaptureControl { .map_or(true, |thread_handle| thread_handle.is_finished()) } + /// Get The Halt Handle Used To Pause The Capture Thread + #[must_use] + pub fn into_thread_handle(self) -> JoinHandle>> { + self.thread_handle.unwrap() + } + + /// Get The Halt Handle Used To Pause The Capture Thread + #[must_use] + pub fn halt_handle(&self) -> Arc { + self.halt_handle.clone() + } + + /// Get The Callback Struct Used To Call Struct Methods Directly + #[must_use] + pub fn callback(&self) -> Arc> { + self.callback.clone() + } + /// Wait Until The Capturing Thread Stops pub fn wait(mut self) -> Result<(), Box> { if let Some(thread_handle) = self.thread_handle.take() { @@ -77,6 +101,8 @@ impl CaptureControl { return Err(Box::new(CaptureControlError::FailedToJoinThread)); } } + } else { + return Err(Box::new(CaptureControlError::ThreadHandleIsTaken)); } Ok(()) @@ -116,6 +142,8 @@ impl CaptureControl { return Err(Box::new(CaptureControlError::FailedToJoinThread)); } } + } else { + return Err(Box::new(CaptureControlError::ThreadHandleIsTaken)); } Ok(()) @@ -150,10 +178,10 @@ pub trait WindowsCaptureHandler: Sized { // Start Capture info!("Starting Capture Thread"); - let trigger = Self::new(settings.flags)?; + let callback = Arc::new(Mutex::new(Self::new(settings.flags)?)); let mut capture = GraphicsCaptureApi::new( settings.item, - trigger, + callback, settings.capture_cursor, settings.draw_border, settings.color_format, @@ -213,12 +241,13 @@ pub trait WindowsCaptureHandler: Sized { /// Starts The Capture Without Taking Control Of The Current Thread fn start_free_threaded( settings: WindowsCaptureSettings, - ) -> Result> + ) -> Result, Box> where Self: Send + 'static, ::Flags: Send, { - let (sender, receiver) = mpsc::channel::>(); + let (halt_sender, halt_receiver) = mpsc::channel::>(); + let (callback_sender, callback_receiver) = mpsc::channel::>>(); let thread_handle = thread::spawn(move || -> Result<(), Box> { // Initialize WinRT @@ -236,10 +265,10 @@ pub trait WindowsCaptureHandler: Sized { // Start Capture info!("Starting Capture Thread"); - let trigger = Self::new(settings.flags)?; + let callback = Arc::new(Mutex::new(Self::new(settings.flags)?)); let mut capture = GraphicsCaptureApi::new( settings.item, - trigger, + callback.clone(), settings.capture_cursor, settings.draw_border, settings.color_format, @@ -249,7 +278,11 @@ pub trait WindowsCaptureHandler: Sized { // Send Halt Handle trace!("Sending Halt Handle"); let halt_handle = capture.halt_handle(); - sender.send(halt_handle)?; + halt_sender.send(halt_handle)?; + + // Send Callback + trace!("Sending Callback"); + callback_sender.send(callback)?; // Debug Thread ID debug!("Thread ID: {}", unsafe { GetCurrentThreadId() }); @@ -301,7 +334,7 @@ pub trait WindowsCaptureHandler: Sized { Ok(()) }); - let halt_handle = match receiver.recv() { + let halt_handle = match halt_receiver.recv() { Ok(halt_handle) => halt_handle, Err(_) => match thread_handle.join() { Ok(result) => return Err(result.err().unwrap()), @@ -311,7 +344,17 @@ pub trait WindowsCaptureHandler: Sized { }, }; - Ok(CaptureControl::new(thread_handle, halt_handle)) + let callback = match callback_receiver.recv() { + Ok(callback_handle) => callback_handle, + Err(_) => match thread_handle.join() { + Ok(result) => return Err(result.err().unwrap()), + Err(_) => { + return Err(Box::new(CaptureControlError::FailedToJoinThread)); + } + }, + }; + + Ok(CaptureControl::new(thread_handle, halt_handle, callback)) } /// Function That Will Be Called To Create The Struct The Flags Can Be diff --git a/src/graphics_capture_api.rs b/src/graphics_capture_api.rs index 2925c6c..8b4ab47 100644 --- a/src/graphics_capture_api.rs +++ b/src/graphics_capture_api.rs @@ -75,7 +75,7 @@ pub struct GraphicsCaptureApi { _d3d_device_context: ID3D11DeviceContext, frame_pool: Option>, session: Option, - closed: Arc, + halt: Arc, active: bool, capture_cursor: Option, draw_border: Option, @@ -85,7 +85,7 @@ impl GraphicsCaptureApi { /// Create A New Graphics Capture Api Struct pub fn new( item: GraphicsCaptureItem, - callback: T, + callback: Arc>, capture_cursor: Option, draw_border: Option, color_format: ColorFormat, @@ -123,21 +123,18 @@ impl GraphicsCaptureApi { trace!("Preallocating Memory"); let mut buffer = vec![0u8; 3840 * 2160 * 4]; - // Trigger Struct - let callback = Arc::new(Mutex::new(callback)); - // Indicates If The Capture Is Closed - let closed = Arc::new(AtomicBool::new(false)); + let halt = Arc::new(AtomicBool::new(false)); // Set Capture Session Closed Event item.Closed( &TypedEventHandler::::new({ // Init let callback_closed = callback.clone(); - let closed_item = closed.clone(); + let halt_closed = halt.clone(); move |_, _| { - closed_item.store(true, atomic::Ordering::Relaxed); + halt_closed.store(true, atomic::Ordering::Relaxed); // To Stop Messge Loop unsafe { PostQuitMessage(0) }; @@ -159,7 +156,7 @@ impl GraphicsCaptureApi { &TypedEventHandler::::new({ // Init let frame_pool_recreate = frame_pool.clone(); - let closed_frame_pool = closed.clone(); + let halt_frame_pool = halt.clone(); let d3d_device_frame_pool = d3d_device.clone(); let context = d3d_device_context.clone(); @@ -169,7 +166,7 @@ impl GraphicsCaptureApi { move |frame, _| { // Return Early If The Capture Is Closed - if closed_frame_pool.load(atomic::Ordering::Relaxed) { + if halt_frame_pool.load(atomic::Ordering::Relaxed) { return Ok(()); } @@ -245,7 +242,7 @@ impl GraphicsCaptureApi { .replace(Some(result)) .expect("Failed To Replace RESULT"); - closed_frame_pool.store(true, atomic::Ordering::Relaxed); + halt_frame_pool.store(true, atomic::Ordering::Relaxed); // To Stop Messge Loop unsafe { PostQuitMessage(0) }; @@ -263,7 +260,7 @@ impl GraphicsCaptureApi { _d3d_device_context: d3d_device_context, frame_pool: Some(frame_pool), session: Some(session), - closed, + halt, active: false, capture_cursor, draw_border, @@ -328,7 +325,7 @@ impl GraphicsCaptureApi { /// Get Halt Handle #[must_use] pub fn halt_handle(&self) -> Arc { - self.closed.clone() + self.halt.clone() } /// Check If Windows Graphics Capture Api Is Supported diff --git a/windows-capture-python/Cargo.toml b/windows-capture-python/Cargo.toml index 6619cc8..e024b49 100644 --- a/windows-capture-python/Cargo.toml +++ b/windows-capture-python/Cargo.toml @@ -4,9 +4,8 @@ version = "1.0.32" authors = ["NiiightmareXD"] edition = "2021" description = "Fastest Windows Screen Capture Library For Python 🔥" -documentation = "https://docs.rs/windows-capture" readme = "README.md" -repository = "https://github.com/NiiightmareXD/windows-capture" +repository = "https://github.com/NiiightmareXD/windows-capture/tree/main/windows-capture-python" license = "MIT" keywords = ["screen", "capture", "screenshot", "graphics", "windows"] categories = [ diff --git a/windows-capture-python/README.md b/windows-capture-python/README.md index 0ffd28e..3c3637e 100644 --- a/windows-capture-python/README.md +++ b/windows-capture-python/README.md @@ -27,9 +27,9 @@ from windows_capture import WindowsCapture, Frame, InternalCaptureControl # Every Error From on_closed and on_frame_arrived Will End Up Here capture = WindowsCapture( - capture_cursor=True, - draw_border=False, - monitor_index=0, + capture_cursor=None, + draw_border=None, + monitor_index=None, window_name=None, ) diff --git a/windows-capture-python/rust-toolchain.toml b/windows-capture-python/rust-toolchain.toml deleted file mode 100644 index fa60a1c..0000000 --- a/windows-capture-python/rust-toolchain.toml +++ /dev/null @@ -1,3 +0,0 @@ -[toolchain] -channel = "stable" -targets = ["x86_64-pc-windows-msvc"] diff --git a/windows-capture-python/rustfmt.toml b/windows-capture-python/rustfmt.toml deleted file mode 100644 index b34697f..0000000 --- a/windows-capture-python/rustfmt.toml +++ /dev/null @@ -1,12 +0,0 @@ -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 c09176c..b71bf65 100644 --- a/windows-capture-python/src/lib.rs +++ b/windows-capture-python/src/lib.rs @@ -7,10 +7,18 @@ #![warn(clippy::cargo)] #![allow(clippy::redundant_pub_crate)] -use std::{error::Error, sync::Arc}; +use std::{ + error::Error, + os::windows::io::AsRawHandle, + sync::{ + atomic::{self, AtomicBool}, + Arc, + }, + thread::JoinHandle, +}; use ::windows_capture::{ - capture::{CaptureControl, WindowsCaptureHandler}, + capture::{CaptureControlError, WindowsCaptureHandler}, frame::Frame, graphics_capture_api::InternalCaptureControl, monitor::Monitor, @@ -18,6 +26,11 @@ use ::windows_capture::{ window::Window, }; use pyo3::{exceptions::PyException, prelude::*, types::PyList}; +use windows::Win32::{ + Foundation::{HANDLE, LPARAM, WPARAM}, + System::Threading::GetThreadId, + UI::WindowsAndMessaging::{PostThreadMessageW, WM_QUIT}, +}; /// Fastest Windows Screen Capture Library For Python 🔥. #[pymodule] @@ -30,11 +43,11 @@ fn windows_capture(_py: Python, m: &PyModule) -> PyResult<()> { /// Internal Struct Used To Handle Free Threaded Start #[pyclass] pub struct NativeCaptureControl { - capture_control: Option, + capture_control: Option, } impl NativeCaptureControl { - const fn new(capture_control: CaptureControl) -> Self { + const fn new(capture_control: NoGenericCaptureControl) -> Self { Self { capture_control: Some(capture_control), } @@ -93,13 +106,94 @@ impl NativeCaptureControl { } } +/// Because The Default Rust Capture Control Contains Generic That Is +/// Unsupported By Python +pub struct NoGenericCaptureControl { + thread_handle: Option>>>, + halt_handle: Arc, +} + +impl NoGenericCaptureControl { + #[must_use] + pub fn new( + thread_handle: JoinHandle>>, + halt_handle: Arc, + ) -> Self { + Self { + thread_handle: Some(thread_handle), + halt_handle, + } + } + + #[must_use] + pub fn is_finished(&self) -> bool { + self.thread_handle + .as_ref() + .map_or(true, |thread_handle| thread_handle.is_finished()) + } + + pub fn wait(mut self) -> Result<(), Box> { + if let Some(thread_handle) = self.thread_handle.take() { + match thread_handle.join() { + Ok(result) => result?, + Err(_) => { + return Err(Box::new(CaptureControlError::FailedToJoinThread)); + } + } + } else { + return Err(Box::new(CaptureControlError::ThreadHandleIsTaken)); + } + + Ok(()) + } + + /// Gracefully Stop The Capture Thread + pub fn stop(mut self) -> Result<(), Box> { + self.halt_handle.store(true, atomic::Ordering::Relaxed); + + if let Some(thread_handle) = self.thread_handle.take() { + let handle = thread_handle.as_raw_handle(); + let handle = HANDLE(handle as isize); + let therad_id = unsafe { GetThreadId(handle) }; + + loop { + match unsafe { + PostThreadMessageW(therad_id, WM_QUIT, WPARAM::default(), LPARAM::default()) + } { + Ok(_) => break, + Err(e) => { + if thread_handle.is_finished() { + break; + } + + if e.code().0 != -2147023452 { + Err(e)?; + } + } + } + } + + match thread_handle.join() { + Ok(result) => result?, + Err(_) => { + return Err(Box::new(CaptureControlError::FailedToJoinThread)); + } + } + } else { + return Err(Box::new(CaptureControlError::ThreadHandleIsTaken)); + } + + Ok(()) + } +} + /// Internal Struct Used For Windows Capture #[pyclass] pub struct NativeWindowsCapture { on_frame_arrived_callback: Arc, on_closed: Arc, - capture_cursor: bool, - draw_border: bool, + capture_cursor: Option, + draw_border: Option, monitor_index: Option, window_name: Option, } @@ -110,9 +204,9 @@ impl NativeWindowsCapture { pub fn new( on_frame_arrived_callback: PyObject, on_closed: PyObject, - capture_cursor: bool, - draw_border: bool, - monitor_index: Option, + capture_cursor: Option, + draw_border: Option, + mut monitor_index: Option, window_name: Option, ) -> PyResult { if window_name.is_some() && monitor_index.is_some() { @@ -122,9 +216,7 @@ impl NativeWindowsCapture { } 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", - )); + monitor_index = Some(0); } Ok(Self { @@ -151,8 +243,8 @@ impl NativeWindowsCapture { match WindowsCaptureSettings::new( window, - Some(self.capture_cursor), - Some(self.draw_border), + self.capture_cursor, + self.draw_border, ColorFormat::Bgra8, ( self.on_frame_arrived_callback.clone(), @@ -178,8 +270,8 @@ impl NativeWindowsCapture { match WindowsCaptureSettings::new( monitor, - Some(self.capture_cursor), - Some(self.draw_border), + self.capture_cursor, + self.draw_border, ColorFormat::Bgra8, ( self.on_frame_arrived_callback.clone(), @@ -221,8 +313,8 @@ impl NativeWindowsCapture { match WindowsCaptureSettings::new( window, - Some(self.capture_cursor), - Some(self.draw_border), + self.capture_cursor, + self.draw_border, ColorFormat::Bgra8, ( self.on_frame_arrived_callback.clone(), @@ -248,8 +340,8 @@ impl NativeWindowsCapture { match WindowsCaptureSettings::new( monitor, - Some(self.capture_cursor), - Some(self.draw_border), + self.capture_cursor, + self.draw_border, ColorFormat::Bgra8, ( self.on_frame_arrived_callback.clone(), @@ -273,7 +365,11 @@ impl NativeWindowsCapture { ))); } }; - let capture_control = NativeCaptureControl::new(capture_control); + + let halt_handle = capture_control.halt_handle(); + let thread_handle = capture_control.into_thread_handle(); + let no_generic_capture_control = NoGenericCaptureControl::new(thread_handle, halt_handle); + let capture_control = NativeCaptureControl::new(no_generic_capture_control); Ok(capture_control) } diff --git a/windows-capture-python/windows_capture/__init__.py b/windows-capture-python/windows_capture/__init__.py index 2880dd5..8f3ea32 100644 --- a/windows-capture-python/windows_capture/__init__.py +++ b/windows-capture-python/windows_capture/__init__.py @@ -159,9 +159,9 @@ class WindowsCapture: def __init__( self, - capture_cursor: bool = True, - draw_border: bool = False, - monitor_index: int = 0, + capture_cursor: Optional[bool] = True, + draw_border: Optional[bool] = None, + monitor_index: Optional[int] = None, window_name: Optional[str] = None, ) -> None: """ diff --git a/windows-capture-python/windows_capture/windows_capture.cp311-win_amd64.pyd b/windows-capture-python/windows_capture/windows_capture.cp311-win_amd64.pyd index 28756a0..bfd8c57 100644 Binary files a/windows-capture-python/windows_capture/windows_capture.cp311-win_amd64.pyd and b/windows-capture-python/windows_capture/windows_capture.cp311-win_amd64.pyd differ