+ Improved type safety with ParamSpec and TypeVar

Author: caffeine-addictt

src/thread/_types.py

@@ -5,6 +5,7 @@
 """
 
 from typing import Any, Literal, Callable, Union
+from typing_extensions import ParamSpec, TypeVar
 
 
 # Descriptive Types
@@ -27,5 +28,10 @@
 
 
 # Function types
-HookFunction = Callable[[Data_Out], Union[Any, None]]
-TargetFunction = Callable[..., Data_Out]
+_Target_P = ParamSpec('_Target_P')
+_Target_T = TypeVar('_Target_T')
+_Dataset_T = TypeVar('_Dataset_T')
+
+TargetFunction = Callable[_Target_P, _Target_T]
+
+HookFunction = Callable[[_Target_T], Union[Any, None]]

src/thread/thread.py

@@ -19,16 +19,22 @@ class ParallelProcessing: ...
 from .utils.config import Settings
 from .utils.algorithm import chunk_split
 
-from ._types import ThreadStatus, Data_In, Data_Out, Overflow_In, TargetFunction, HookFunction
+from ._types import (
+  ThreadStatus, Data_In, Data_Out, Overflow_In,
+  TargetFunction, _Target_P, _Target_T,
+  DatasetFunction, _Dataset_T,
+  HookFunction
+)
+from typing_extensions import Generic
 from typing import (
-  Any, List,
-  Callable, Optional,
+  Any, List, Unpack,
+  Callable, Optional, Union,
   Mapping, Sequence, Tuple
 )
 
 
 Threads: set['Thread'] = set()
-class Thread(threading.Thread):
+class Thread(threading.Thread, Generic[_Target_P, _Target_T]):
   """
   Wraps python's `threading.Thread` class
   ---------------------------------------
@@ -51,7 +57,7 @@ class Thread(threading.Thread):
 
   def __init__(
     self,
-    target: TargetFunction,
+    target: TargetFunction[_Target_P, _Target_T],
     args: Sequence[Data_In] = (),
     kwargs: Mapping[str, Data_In] = {},
     ignore_errors: Sequence[type[Exception]] = (),
@@ -100,10 +106,10 @@ def __init__(
     )
 
 
-  def _wrap_target(self, target: TargetFunction) -> TargetFunction:
+  def _wrap_target(self, target: TargetFunction[_Target_P, _Target_T]) -> TargetFunction[_Target_P, Union[_Target_T, None]]:
     """Wraps the target function"""
     @wraps(target)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
+    def wrapper(*args: _Target_P.args, **kwargs: _Target_P.kwargs) -> Union[_Target_T, None]:
       self.status = 'Running'
 
       global Threads
@@ -173,7 +179,7 @@ def _run_with_trace(self) -> None:
 
 
   @property
-  def result(self) -> Data_Out:
+  def result(self) -> _Target_T:
     """
     The return value of the thread
     
@@ -208,7 +214,7 @@ def is_alive(self) -> bool:
     return super().is_alive()
 
 
-  def add_hook(self, hook: HookFunction) -> None:
+  def add_hook(self, hook: HookFunction[_Target_T]) -> None:
     """
     Adds a hook to the thread
     -------------------------
@@ -250,7 +256,7 @@ def join(self, timeout: Optional[float] = None) -> bool:
     return not self.is_alive()
 
 
-  def get_return_value(self) -> Data_Out:
+  def get_return_value(self) -> _Target_T:
     """
     Halts the current thread execution until the thread completes
 
@@ -323,7 +329,7 @@ def __init__(self, thread: Thread, progress: float = 0) -> None:
     self.thread = thread
     self.progress = progress
 
-class ParallelProcessing:
+class ParallelProcessing(Generic[_Target_P, _Target_T, _Dataset_T]):
   """
   Multi-Threaded Parallel Processing
   ---------------------------------------
@@ -335,7 +341,7 @@ class ParallelProcessing:
   _completed     : int
 
   status         : ThreadStatus
-  function       : Callable[..., List[Data_Out]]
+  function       : TargetFunction[..., List[_Target_T]]
   dataset        : Sequence[Data_In]
   max_threads    : int
 
@@ -344,8 +350,8 @@ class ParallelProcessing:
 
   def __init__(
     self,
-    function: TargetFunction,
-    dataset: Sequence[Data_In],
+    function: DatasetFunction[_Dataset_T, _Target_T],
+    dataset: Sequence[_Dataset_T],
     max_threads: int = 8,
 
     *overflow_args: Overflow_In,
@@ -385,10 +391,10 @@ def __init__(
 
   def _wrap_function(
     self,
-    function: TargetFunction
-  ) -> Callable[..., List[Data_Out]]:
+    function: TargetFunction[[_Dataset_T], _Target_T]
+  ) -> TargetFunction[..., List[_Target_T]]:
     @wraps(function)
-    def wrapper(index: int, data_chunk: Sequence[Data_In], *args: Any, **kwargs: Any) -> List[Data_Out]:
+    def wrapper(index: int, data_chunk: Sequence[_Dataset_T], *args: _Target_P.args, **kwargs: _Target_P.kwargs) -> List[_Target_T]:
       computed: List[Data_Out] = []
       for i, data_entry in enumerate(data_chunk):
         v = function(data_entry, *args, **kwargs)
@@ -404,7 +410,7 @@ def wrapper(index: int, data_chunk: Sequence[Data_In], *args: Any, **kwargs: Any
 
 
   @property
-  def results(self) -> Data_Out:
+  def results(self) -> List[_Dataset_T]:
     """
     The return value of the threads if completed
     
@@ -436,7 +442,7 @@ def is_alive(self) -> bool:
     return any(entry.thread.is_alive() for entry in self._threads)
 
 
-  def get_return_values(self) -> List[Data_Out]:
+  def get_return_values(self) -> List[_Dataset_T]:
     """
     Halts the current thread execution until the thread completes