Merge PR 18

Author: zh1peng

psyflow/BlockUnit.py

@@ -1,5 +1,11 @@
+"""Block-level trial controller.
+
+Manages condition generation (weighted, balanced, or custom), trial execution
+with lifecycle hooks, and per-block result aggregation.
+"""
+
 import numpy as np
-from typing import Callable, Any, List, Dict, Optional
+from typing import Callable, Any, List, Dict, Optional, overload
 from psychopy import core, logging
 from typing import Union, List, Dict, Literal
 import re
@@ -198,7 +204,15 @@ def add_condition(self, condition_list: List[Any]) -> "BlockUnit":
         self.conditions = condition_list
         return self
 
-    def on_start(self, func: Optional[Callable[['BlockUnit'], None]] = None):
+    @overload
+    def on_start(self, func: None = None) -> Callable[[Callable[['BlockUnit'], None]], 'BlockUnit']:
+        ...
+
+    @overload
+    def on_start(self, func: Callable[['BlockUnit'], None]) -> 'BlockUnit':
+        ...
+
+    def on_start(self, func: Optional[Callable[['BlockUnit'], None]] = None) -> Union['BlockUnit', Callable[[Callable[['BlockUnit'], None]], 'BlockUnit']]:
         """
         Register a function to run at the start of the block.
 
@@ -215,7 +229,15 @@ def decorator(f):
         self._on_start.append(func)
         return self
 
-    def on_end(self, func: Optional[Callable[['BlockUnit'], None]] = None):
+    @overload
+    def on_end(self, func: None = None) -> Callable[[Callable[['BlockUnit'], None]], 'BlockUnit']:
+        ...
+
+    @overload
+    def on_end(self, func: Callable[['BlockUnit'], None]) -> 'BlockUnit':
+        ...
+
+    def on_end(self, func: Optional[Callable[['BlockUnit'], None]] = None) -> Union['BlockUnit', Callable[[Callable[['BlockUnit'], None]], 'BlockUnit']]:
         """
         Register a function to run at the end of the block.
 
@@ -232,7 +254,7 @@ def decorator(f):
         self._on_end.append(func)
         return self
 
-    def run_trial(self, func: Callable, **kwargs):
+    def run_trial(self, func: Callable, **kwargs) -> "BlockUnit":
         """
         Run all trials using a specified trial function.
 
@@ -377,7 +399,7 @@ def match(value: str) -> bool:
             if negate ^ match(str(trial.get(key, '')))
         ]
 
-    def logging_block_info(self):
+    def logging_block_info(self) -> None:
         """
         Log block metadata including ID, index, seed, trial count, and condition distribution.
         """

psyflow/StimBank.py

@@ -1,3 +1,9 @@
+"""Stimulus registry with lazy instantiation.
+
+Supports decorator-based and YAML/dict-based stimulus definitions, batch
+preview, text formatting, and text-to-speech conversion via edge-tts.
+"""
+
 from psychopy.visual import TextStim, Circle, Rect, Polygon, ImageStim, ShapeStim, TextBox2, MovieStim
 from psychopy import event, core
 
@@ -72,7 +78,7 @@ def decorator(func: Callable[[Any], Any]):
             return func
         return decorator
 
-    def preload_all(self):
+    def preload_all(self) -> "StimBank":
         """Instantiate all registered stimuli.
 
         Returns
@@ -214,7 +220,7 @@ def get_selected(self, keys: list[str]) -> Dict[str, Any]:
         """
         return {k: self.get(k) for k in keys}
 
-    def preview_all(self, wait_keys: bool = True):
+    def preview_all(self, wait_keys: bool = True) -> None:
         """
         Preview all registered stimuli one by one.
 
@@ -227,7 +233,7 @@ def preview_all(self, wait_keys: bool = True):
         for i, name in enumerate(keys):
             self._preview(name, wait_keys=wait_keys)
 
-    def preview_group(self, prefix: str, wait_keys: bool = True):
+    def preview_group(self, prefix: str, wait_keys: bool = True) -> None:
         """
         Preview all stimuli that match a name prefix.
 
@@ -244,7 +250,7 @@ def preview_group(self, prefix: str, wait_keys: bool = True):
         for i, name in enumerate(matches):
             self._preview(name, wait_keys=(i == len(matches) - 1))
 
-    def preview_selected(self, keys: list[str], wait_keys: bool = True):
+    def preview_selected(self, keys: list[str], wait_keys: bool = True) -> None:
         """
         Preview selected stimuli by name.
 
@@ -258,29 +264,7 @@ def preview_selected(self, keys: list[str], wait_keys: bool = True):
         for i, name in enumerate(keys):
             self._preview(name, wait_keys=(i == len(keys) - 1))
 
-    # def _preview(self, name: str, wait_keys: bool = True):
-    #     """
-    #     Internal utility to preview a single stimulus.
-
-    #     Parameters
-    #     ----------
-    #     name : str
-    #         Stimulus name.
-    #     wait_keys : bool
-    #         Wait for key press after preview.
-    #     """
-    #     try:
-    #         stim = self.get(name)
-    #         self.win.flip(clearBuffer=True)
-    #         stim.draw()
-    #         self.win.flip()
-    #         print(f"Preview: '{name}'")
-    #         if wait_keys:
-    #             event.waitKeys()
-    #     except Exception as e:
-    #         print(f"[Preview Error] Could not preview '{name}': {e}")
-
-    def _preview(self, name: str, wait_keys: bool = True):
+    def _preview(self, name: str, wait_keys: bool = True) -> None:
         """
         Internal utility to preview a single stimulus (image or sound).
 
@@ -337,7 +321,7 @@ def has(self, name: str) -> bool:
         """
         return name in self._registry
 
-    def describe(self, name: str):
+    def describe(self, name: str) -> None:
         """
         Print accepted arguments for a registered stimulus.
 
@@ -370,7 +354,7 @@ def describe(self, name: str):
             default = "required" if v.default is inspect.Parameter.empty else f"default={v.default!r}"
             print(f"  - {k}: {default}")
 
-    def export_to_yaml(self, path: str):
+    def export_to_yaml(self, path: str) -> None:
         """
         Export YAML-defined stimuli (but not decorator-defined) to file.
 
@@ -382,6 +366,8 @@ def export_to_yaml(self, path: str):
         yaml_defs = {}
         for name, factory in self._registry.items():
             try:
+                # Factories created by add_from_dict() capture their source
+                # dict in a closure. Inspect it to recover the original spec.
                 source = factory.__closure__[0].cell_contents
                 if not isinstance(source, dict):
                     continue
@@ -393,7 +379,7 @@ def export_to_yaml(self, path: str):
             yaml.dump(yaml_defs, f)
         print(f"[OK] Exported {len(yaml_defs)} YAML stimuli to {path}")
 
-    def make_factory(self, cls, base_kwargs: dict, name: str):
+    def make_factory(self, cls: type, base_kwargs: dict, name: str) -> Callable:
         """
         Create a factory function for a given stimulus class.
 
@@ -426,7 +412,7 @@ def _factory(win, **override_kwargs):
                 raise ValueError(f"[StimBank] Failed to build '{name}': {e}")
         return _factory
 
-    def add_from_dict(self, named_specs: Optional[dict] = None, **kwargs):
+    def add_from_dict(self, named_specs: Optional[dict] = None, **kwargs) -> "StimBank":
         """
         Add stimuli from a dictionary or keyword-based specifications.
 
@@ -452,7 +438,7 @@ def add_from_dict(self, named_specs: Optional[dict] = None, **kwargs):
             self._registry[name] = self.make_factory(stim_class, kwargs, name)
         return self
 
-    def validate_dict(self, config: dict, strict: bool = False):
+    def validate_dict(self, config: dict, strict: bool = False) -> None:
         """
         Validate a dictionary of stimulus definitions.
 
@@ -506,7 +492,7 @@ def validate_dict(self, config: dict, strict: bool = False):
     def convert_to_voice(self,
                          keys: list[str] | str,
                          overwrite: bool = False,
-                         voice: str = "zh-CN-YunyangNeural"):
+                         voice: str = "zh-CN-YunyangNeural") -> "StimBank":
         """
         Convert specified TextStim/TextBox2 stimuli to speech (MP3) and register them
         as new Sound stimuli in this StimBank.
@@ -578,7 +564,7 @@ def add_voice(self,
                   stim_label: str,
                   text: str,
                   overwrite: bool = False,
-                  voice: str = "zh-CN-XiaoxiaoNeural"):
+                  voice: str = "zh-CN-XiaoxiaoNeural") -> "StimBank":
         """
         Convert arbitrary text to speech (MP3) and register it as a new Sound stimulus.
 

psyflow/StimUnit.py

@@ -1,3 +1,10 @@
+"""Trial-level stimulus executor.
+
+Encapsulates stimulus presentation, response capture, event triggers,
+timing control, and lifecycle hooks.  Adapts automatically to simulation
+mode via :class:`~psyflow.sim.adapter.ResponderAdapter`.
+"""
+
 from psychopy import core, visual, logging, sound
 from psychopy.hardware.keyboard import Keyboard
 from typing import Callable, Optional, List, Dict, Any, Sequence, TypeAlias, Union

psyflow/SubInfo.py

@@ -1,3 +1,10 @@
+"""Participant information dialog.
+
+Presents a configurable PsychoPy GUI dialog to collect and validate
+participant metadata (subject ID, demographics, etc.) with optional
+localization support.
+"""
+
 from psychopy import gui
 
 class SubInfo:

psyflow/TaskSettings.py

@@ -1,3 +1,10 @@
+"""Experiment configuration container.
+
+Holds window display, block/trial structure, seeding strategy, and per-subject
+output paths.  Instantiate directly or via :meth:`TaskSettings.from_dict` with
+a YAML-loaded dictionary.
+"""
+
 from dataclasses import dataclass, field
 from typing import List, Optional, Any, Dict
 from math import ceil
@@ -70,7 +77,7 @@ def __post_init__(self):
         if self.seed_mode == 'same_across_sub' and all(seed is None for seed in self.block_seed):
             self.set_block_seed(self.overall_seed)
 
-    def set_block_seed(self, seed_base: Optional[int]):
+    def set_block_seed(self, seed_base: Optional[int]) -> None:
         """
         Generate a list of per-block seeds from a base seed.
 
@@ -142,7 +149,7 @@ def resolve_condition_weights(self) -> list[float] | None:
             raise ValueError(f"condition_weights sum must be > 0, got {weights}")
         return weights
 
-    def add_subinfo(self, subinfo: Dict[str, Any]):
+    def add_subinfo(self, subinfo: Dict[str, Any]) -> None:
         """
         Add subject-specific information and set seed/file names accordingly.
 
@@ -187,14 +194,14 @@ def add_subinfo(self, subinfo: Dict[str, Any]):
         self.res_file = os.path.join(self.save_path, f"{basename}.csv")
         self.json_file = os.path.join(self.save_path, f"{basename}.json")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Return a clean string representation of the current TaskSettings.
         """
         base = {k: v for k, v in self.__dict__.items() if not k.startswith('_')}
         return f"{self.__class__.__name__}({base})"
 
-    def save_to_json(self):
+    def save_to_json(self) -> None:
         """
         Save the current TaskSettings instance to a JSON file.
         """
@@ -218,7 +225,7 @@ def save_to_json(self):
 
 
     @classmethod
-    def from_dict(cls, config: dict):
+    def from_dict(cls, config: dict) -> "TaskSettings":
         """
         Create a TaskSettings instance from a flat dictionary.
 

psyflow/utils/display.py

@@ -3,7 +3,7 @@
 from psychopy import core, visual
 
 
-def count_down(win, seconds=3, **stim_kwargs):
+def count_down(win: "visual.Window", seconds: int = 3, **stim_kwargs) -> None:
     """Display a frame-accurate countdown using TextStim."""
     cd_clock = core.Clock()
     for i in reversed(range(1, seconds + 1)):

psyflow/utils/ports.py

@@ -1,7 +1,7 @@
 """Serial port helper utilities."""
 
 
-def show_ports():
+def show_ports() -> None:
     """List all available serial ports."""
     import serial.tools.list_ports
 

psyflow/utils/templates.py

@@ -4,7 +4,7 @@
 import importlib.resources as pkg_res
 
 
-def taps(task_name: str, template: str = "cookiecutter-psyflow"):
+def taps(task_name: str, template: str = "cookiecutter-psyflow") -> str:
     """Generate a task skeleton using the bundled template."""
     tmpl_dir = pkg_res.files("psyflow.templates") / template
     cookiecutter(

psyflow/utils/trials.py

@@ -1,3 +1,5 @@
+"""Trial ID generation and deadline resolution utilities."""
+
 from typing import Any
 
 _SESSION_TRIAL_COUNTER = 0
@@ -8,7 +10,7 @@ def next_trial_id() -> int:
     _SESSION_TRIAL_COUNTER += 1
     return _SESSION_TRIAL_COUNTER
 
-def reset_trial_counter(start_at: int = 0):
+def reset_trial_counter(start_at: int = 0) -> None:
     """Reset the global trial counter."""
     global _SESSION_TRIAL_COUNTER
     _SESSION_TRIAL_COUNTER = start_at

psyflow/utils/voices.py

@@ -17,7 +17,7 @@ async def _list_supported_voices_async(filter_lang: Optional[str] = None):
 def list_supported_voices(
     filter_lang: Optional[str] = None,
     human_readable: bool = False,
-):
+) -> list[dict] | None:
     """Query available edge-tts voices."""
     voices = asyncio.run(_list_supported_voices_async(filter_lang))
     if not human_readable: