feat(context): Add context subclass feature (#114)

* feat(context): Start working on contexts feature

* Fix clippy error by moving args to a struct

* Move structs below the run_worker func

* Start working on context in the worker

* Fix clippy error

* Remove dispatcher pool, was no different without it

* Add tests for tasks

* Remove commented lines

* Finish context subclasses feature

* Fix clippy errors

* Add process exit on task registry init failure

* Add thread storage to the context

* Add run-worker wrapper script

* Fix unnecessary check in context decorator

* Update thread storage to ContextVar instead of thread.local()

* Add task metadata feat

* Finish

* Update cargo tests

* Fix tests

* Fix tests

* Fix tests

* Fix tests

* Fix tests once again

* Fix tests

* Move tests to bash script

* Fix run-tests

* Fix run-tests

* Fix tests

* Update tests

* Revert "Update tests"

This reverts commit 7f1bd54.

* Remove uv installation

Author: CCXLV

.github/actions/tests/action.yml

@@ -13,12 +13,9 @@ runs:
       with:
         python-version: ${{ inputs.python-version }}
 
-    - name: Install uv
-      uses: astral-sh/setup-uv@v7
-
     - name: Install dependencies
       shell: bash -euxo pipefail {0}
-      run: uv sync --extra tests
+      run: pip install -e .[tests]
 
     - name: Set up Rust
       uses: dtolnay/rust-toolchain@stable
@@ -29,5 +26,5 @@ runs:
     - name: Run Tests
       shell: bash -euxo pipefail {0}
       run: |
-        uv run pytest tests/
+        pytest tests/
         cargo test --workspace --locked

.python-version

@@ -0,0 +1 @@
+3.13.11

crates/fluxqueue-worker/Cargo.toml

@@ -12,6 +12,7 @@ name = "fluxqueue-worker"
 path = "src/main.rs"
 
 [dependencies]
+pyo3 = { version = "0.27.2", features = ["auto-initialize"] }
 pyo3-async-runtimes = { version = "0.27.0", features = ["tokio-runtime"] }
 uuid = { version = "1.20.0", features = ["v4"] }
 redis = "0.32.7"
@@ -22,12 +23,15 @@ tokio = { version = "1.49.0", features = [
   "signal",
   "time",
 ] }
-pyo3 = { version = "0.27.2", features = ["auto-initialize"] }
 rmp-serde = "1.3.1"
 rmpv = { version = "1.3.1", features = ["with-serde"] }
 serde = { version = "1.0.228", features = ["derive"] }
 tracing = "0.1.44"
-tracing-subscriber = { version = "0.3.22", features = ["env-filter", "fmt", "time"] }
+tracing-subscriber = { version = "0.3.22", features = [
+  "env-filter",
+  "fmt",
+  "time",
+] }
 time = { version = "0.3.47", features = ["local-offset", "macros"] }
 clap = { version = "4.5.56", features = ["derive", "env"] }
 anyhow = "1.0.100"

crates/fluxqueue-worker/scripts/get_functions.py

@@ -0,0 +1,59 @@
+import importlib
+import inspect
+import sys
+from pathlib import Path
+from typing import get_type_hints
+
+from fluxqueue import Context
+
+
+def get_registry(module_path: str, queue: str, module_dir: str | None = None):
+    if module_dir:
+        module_dir_path = Path(module_dir).resolve()
+        if str(module_dir_path) not in sys.path:
+            sys.path.insert(0, str(module_dir_path))
+
+    module = importlib.import_module(module_path)
+    registry = {"tasks": {}, "contexts": {}}
+    for _name, obj in inspect.getmembers(module):
+        if inspect.isfunction(obj):
+            task_name = getattr(obj, "task_name", None)
+            task_queue = getattr(obj, "queue", None)
+            if not task_queue or task_queue != queue:
+                continue
+
+            if registry["tasks"].get(task_name):
+                raise ValueError(f"Task '{task_name}' is duplicated")
+
+            original_func = getattr(obj, "__wrapped__", obj)
+
+            hints = get_type_hints(original_func)
+            sig = inspect.signature(original_func)
+            context_params = {
+                name: hints[name]
+                for name in sig.parameters
+                if name in hints
+                and isinstance(hints[name], type)
+                and issubclass(hints[name], Context)
+            }
+            if not context_params:
+                context_name = None
+            else:
+                context = context_params[next(iter(context_params))]
+                context_name = getattr(context, "__fluxqueue_context__", None)
+
+            registry["tasks"][task_name] = {
+                "func": original_func,
+                "context_name": context_name,
+            }
+        elif inspect.isclass(obj):
+            if not issubclass(obj, Context):
+                continue
+
+            context_name = getattr(obj, "__fluxqueue_context__", None)
+            if registry["contexts"].get(context_name):
+                raise ValueError(f"Context '{context_name}' is duplicated")
+
+            registry["contexts"][context_name] = obj
+
+    return registry

crates/fluxqueue-worker/scripts/get_registry.py

@@ -34,13 +34,15 @@ pub fn initial_logs(
     concurrency: usize,
     redis_url: &str,
     tasks_module_path: &str,
-    tasks: &Vec<&String>,
+    tasks: Vec<String>,
+    contexts: Vec<String>,
 ) {
     info!("Queue: {}", queue_name);
     info!("Concurrency: {}", concurrency);
     info!("Redis: {}", redis_url);
     info!("Tasks module: {}", tasks_module_path);
     info!("Tasks found: {:?}", tasks);
+    info!("Contexts found: {:?}", contexts);
     info!("Starting up the executors...");
 }
 

crates/fluxqueue-worker/src/logger.rs

@@ -33,7 +33,7 @@ impl RedisClient {
         Ok(())
     }
 
-    pub async fn set_executors_heartbeat(&self, executor_ids: Arc<Vec<Arc<str>>>) -> Result<()> {
+    pub async fn set_executors_heartbeat(&self, executor_ids: Arc<Vec<Arc<String>>>) -> Result<()> {
         for id in executor_ids.iter() {
             self.set_executor_heartbeat(id).await?;
         }
@@ -60,7 +60,7 @@ impl RedisClient {
     pub async fn cleanup_executors_registry(
         &self,
         queue_name: &str,
-        ids: Arc<Vec<Arc<str>>>,
+        ids: Arc<Vec<Arc<String>>>,
     ) -> Result<()> {
         let mut conn = self.redis_pool.get().await?;
         let executors_key = keys::get_executors_key(queue_name);