From 48397c06fadd6b3f86a20ef9ea15f931c23ca774 Mon Sep 17 00:00:00 2001
From: Leo dev <klesti.selimay@icloud.com>
Date: Mon, 4 Aug 2025 04:45:55 -0500
Subject: [PATCH] Added ExtensionContext #26

---
 src/extensions/input_handling.rs | 11 +++----
 src/extensions/mod.rs            | 53 +++++++++++++++++++++++++++++---
 src/extensions/tick.rs           | 13 ++++----
 src/extensions/velocity.rs       |  6 ++--
 src/lib.rs                       |  7 +++--
 5 files changed, 68 insertions(+), 22 deletions(-)

diff --git a/src/extensions/input_handling.rs b/src/extensions/input_handling.rs
index 227d526..8939812 100644
--- a/src/extensions/input_handling.rs
+++ b/src/extensions/input_handling.rs
@@ -1,20 +1,19 @@
-use crate::extensions::Extension;
+use crate::extensions::{Extension, ExtensionContext};
 
 pub struct InputExtension;
 
 impl Extension for InputExtension {
-    fn init(&mut self, screen: std::sync::Arc<crate::Screen>) {
+    fn init(&mut self, ctx: &ExtensionContext) {
+        let ctx = ctx.clone();
         crossterm::terminal::enable_raw_mode().unwrap();
         std::thread::spawn(move || loop {
             if let Ok(e) = crossterm::event::read() {
-                for widget in screen.widgets.lock().unwrap().iter() {
-                    widget.event(&e);
-                }
+                ctx.event(&e);
             }
         });
     }
 
-    fn on_close(&mut self, _screen: std::sync::Arc<crate::Screen>) {
+    fn on_close(&mut self) {
         crossterm::terminal::disable_raw_mode().unwrap();
     }
 }
diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs
index d027629..ee53bf9 100644
--- a/src/extensions/mod.rs
+++ b/src/extensions/mod.rs
@@ -11,7 +11,7 @@ pub use velocity::*;
 use std::{
     any::{type_name, Any},
     fmt::Debug,
-    sync::{Arc, Mutex},
+    sync::{Arc, Mutex, MutexGuard},
 };
 
 use crate::{
@@ -21,15 +21,25 @@ use crate::{
 };
 
 pub trait Extension {
-    fn init(&mut self, _screen: Arc<Screen>) {}
-    fn on_close(&mut self, _screen: Arc<Screen>) {}
-    fn render_widget(&mut self, _scope: &mut RenderScope, _widget: &Arc<Widget>) {}
+    fn init(&mut self, _ctx: &ExtensionContext) {}
+    fn on_close(&mut self) {}
+    fn render_widget(
+        &mut self,
+        _scope: &mut RenderScope,
+        _widget: &Arc<Widget>,
+    ) {
+    }
 }
 
 pub trait Event: Send + Sync {
     fn as_any(&self) -> &dyn Any;
 }
 
+#[derive(Clone)]
+pub struct ExtensionContext {
+    screen: Arc<Screen>,
+}
+
 #[derive(Clone)]
 pub struct Handler<E: Event>(Arc<Mutex<dyn FnMut(&Arc<Widget>, &E) + Send + Sync>>);
 
@@ -64,3 +74,38 @@ impl<'a> dyn Event + 'a {
         self.as_any().downcast_ref()
     }
 }
+
+impl ExtensionContext {
+    pub fn new(screen: Arc<Screen>) -> Self {
+        Self { screen }
+    }
+
+    pub fn event<E: Event + Clone + 'static>(&self, e: &E) {
+        for widget in self.screen.widgets.lock().unwrap().iter() {
+            widget.event(e);
+        }
+    }
+
+    pub fn get_widgets(&self) -> MutexGuard<Vec<Arc<Widget>>> {
+        self.screen.widgets.lock().unwrap()
+    }
+
+    pub fn iter_components<C: Component + 'static + Clone, F: FnMut(&Arc<Widget>, Option<C>)>(
+        &self,
+        mut iterator: F,
+    ) {
+        for widget in self.screen.widgets.lock().unwrap().iter() {
+            iterator(widget, widget.get());
+        }
+    }
+
+    pub fn get_components<C: Component + 'static + Clone>(&self) -> Vec<C> {
+        let mut components = Vec::new();
+        for widget in self.screen.widgets.lock().unwrap().iter() {
+            if let Some(c) = widget.get() {
+                components.push(c);
+            }
+        }
+        components
+    }
+}
diff --git a/src/extensions/tick.rs b/src/extensions/tick.rs
index 354814d..f48ce61 100644
--- a/src/extensions/tick.rs
+++ b/src/extensions/tick.rs
@@ -1,22 +1,21 @@
-use std::{fmt::Debug, sync::Arc};
+use std::fmt::Debug;
 
-use crate::extensions::Extension;
-use crate::{event, Screen};
+use crate::event;
+use crate::extensions::{Extension, ExtensionContext};
 
 pub struct TickExtension(pub u16);
 
 event!(TickEvent(pub u32));
 
 impl Extension for TickExtension {
-    fn init(&mut self, screen: Arc<Screen>) {
+    fn init(&mut self, ctx: &ExtensionContext) {
+        let ctx = ctx.clone();
         let rate_dur = 1000 / self.0 as u64;
         std::thread::spawn({
             move || {
                 let mut tick = 0;
                 loop {
-                    for w in screen.widgets.lock().unwrap().iter() {
-                        w.event(&TickEvent(tick));
-                    }
+                    ctx.event(&TickEvent(tick));
                     tick += 1;
                     std::thread::sleep(std::time::Duration::from_millis(rate_dur));
                 }
diff --git a/src/extensions/velocity.rs b/src/extensions/velocity.rs
index 74aabee..799f8a0 100644
--- a/src/extensions/velocity.rs
+++ b/src/extensions/velocity.rs
@@ -5,7 +5,6 @@ use crate::{
     extensions::Extension,
     style::{Position, Transform},
     widget::Widget,
-    Screen,
 };
 
 pub struct VelocityExtension;
@@ -43,12 +42,13 @@ impl VelocityExtension {
 }
 
 impl Extension for VelocityExtension {
-    fn init(&mut self, screen: Arc<Screen>) {
+    fn init(&mut self, ctx: &super::ExtensionContext) {
+        let ctx = ctx.clone();
         std::thread::spawn({
             move || {
                 let mut tick = 0;
                 loop {
-                    for widget in screen.widgets.lock().unwrap().iter() {
+                    for widget in ctx.get_widgets().iter() {
                         Self::apply_velocity_xy(tick, widget);
                     }
                     if tick > 1000 {
diff --git a/src/lib.rs b/src/lib.rs
index 135fbc5..c82b1cb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -33,6 +33,7 @@ use std::sync::{Arc, Mutex};
 
 use crate::{
     extensions::{Extension, Handler},
+    prelude::ExtensionContext,
     render_scope::RenderScope,
     widget::{BoxedElement, DynWidget, Element, StaticWidget, Widget, WidgetLoad},
 };
@@ -158,8 +159,10 @@ impl Screen {
     ///
     /// This method blocks and repeatedly renders the screen at a fixed interval.
     pub fn run(self: &Arc<Self>) -> std::io::Result<()> {
+        let ctx = ExtensionContext::new(self.clone());
+
         for ext in self.extensions.lock().unwrap().iter() {
-            ext.lock().unwrap().init(self.clone());
+            ext.lock().unwrap().init(&ctx);
         }
 
         utils::hide_cursor()?;
@@ -230,7 +233,7 @@ impl Screen {
         utils::clear().unwrap();
 
         for ext in self.extensions.lock().unwrap().iter() {
-            ext.lock().unwrap().on_close(self.clone());
+            ext.lock().unwrap().on_close();
         }
     }
 }