Refactor PyOsmo core to simplify error handling, improve registries, and enhance algorithm initialization.

- Centralize error handling for "failure_in_suite" to reuse "failure_in_test."
- Consolidate registries (algorithms, end conditions, error strategies) into a unified `_registries` structure for maintainability.
- Add `_ensure_initialized` method in algorithm base to validate initialization.
- Simplify history counts and logs using comprehensions.
- Standardize string styles and simplify Playwright-based model scripts.
- Remove redundant error handling methods.

Author: OPpuolitaival

examples/ai_web_agent/example_output/todo_app_model.py

@@ -22,16 +22,16 @@ def __init__(self, page: Page, url: str):
     def before_test(self):
         """Navigate to app and reset state."""
         self.page.goto(self.url)
-        self.page.wait_for_selector(".new-todo")
+        self.page.wait_for_selector('.new-todo')
         self.todo_count = 0
         self.completed_count = 0
 
     # --- Add a todo item ---
 
     def step_add_todo(self):
-        input_field = self.page.locator(".new-todo")
-        input_field.fill(f"Task {self.todo_count + 1}")
-        input_field.press("Enter")
+        input_field = self.page.locator('.new-todo')
+        input_field.fill(f'Task {self.todo_count + 1}')
+        input_field.press('Enter')
         self.todo_count += 1
 
     def guard_add_todo(self):
@@ -43,7 +43,7 @@ def weight_add_todo(self):
     # --- Toggle a todo as complete ---
 
     def step_toggle_todo(self):
-        items = self.page.locator(".todo-list li:not(.completed) .toggle")
+        items = self.page.locator('.todo-list li:not(.completed) .toggle')
         items.first.click()
         self.completed_count += 1
 
@@ -54,9 +54,9 @@ def guard_toggle_todo(self):
     # --- Delete a todo item ---
 
     def step_delete_todo(self):
-        item = self.page.locator(".todo-list li").first
+        item = self.page.locator('.todo-list li').first
         item.hover()
-        item.locator(".destroy").click()
+        item.locator('.destroy').click()
         self.todo_count -= 1
 
     def guard_delete_todo(self):
@@ -92,18 +92,18 @@ def after(self):
         """Verify todo count display matches model state."""
         active_count = self.todo_count - self.completed_count
         if active_count > 0:
-            count_text = self.page.locator(".todo-count").text_content()
+            count_text = self.page.locator('.todo-count').text_content()
             assert str(active_count) in count_text
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     from playwright.sync_api import sync_playwright
 
     with sync_playwright() as p:
         browser = p.chromium.launch(headless=False)
         page = browser.new_page()
 
-        model = TodoAppModel(page=page, url="https://todomvc.com/examples/react/dist/")
+        model = TodoAppModel(page=page, url='https://todomvc.com/examples/react/dist/')
 
         osmo = Osmo(model)
         osmo.test_end_condition = Length(30)

examples/ai_web_agent/generate_model.py

@@ -12,7 +12,6 @@
 import argparse
 import asyncio
 import sys
-from pathlib import Path
 
 from prompt_template import PYOSMO_MODEL_REFERENCE
 
@@ -21,17 +20,17 @@ async def generate_model(url: str, output_path: str) -> None:
     try:
         from claude_agent_sdk import Agent, AgentConfig, MCPServer
     except ImportError:
-        print("Error: claude-agent-sdk is required. Install with: pip install claude-agent-sdk")
+        print('Error: claude-agent-sdk is required. Install with: pip install claude-agent-sdk')
         sys.exit(1)
 
     playwright_mcp = MCPServer(
-        name="playwright",
-        command="npx",
-        args=["@playwright/mcp@latest"],
+        name='playwright',
+        command='npx',
+        args=['@playwright/mcp@latest'],
     )
 
     agent = Agent(
-        model="claude-sonnet-4-5-20250929",
+        model='claude-sonnet-4-5-20250929',
         config=AgentConfig(
             system_prompt=PYOSMO_MODEL_REFERENCE,
             mcp_servers=[playwright_mcp],
@@ -51,20 +50,22 @@ async def generate_model(url: str, output_path: str) -> None:
 Save the model to: {output_path}
 """
 
-    print(f"Exploring {url} and generating model...")
+    print(f'Exploring {url} and generating model...')
     result = await agent.run(user_prompt)
-    print(f"Agent completed. Model saved to {output_path}")
-    print(f"Result: {result}")
+    print(f'Agent completed. Model saved to {output_path}')
+    print(f'Result: {result}')
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description="Generate a PyOsmo model from a web application")
-    parser.add_argument("url", help="URL of the web application to model")
-    parser.add_argument("--output", "-o", default="generated_model.py", help="Output file path (default: generated_model.py)")
+    parser = argparse.ArgumentParser(description='Generate a PyOsmo model from a web application')
+    parser.add_argument('url', help='URL of the web application to model')
+    parser.add_argument(
+        '--output', '-o', default='generated_model.py', help='Output file path (default: generated_model.py)'
+    )
     args = parser.parse_args()
 
     asyncio.run(generate_model(args.url, args.output))
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()

examples/ai_web_agent/refine_model.py

@@ -32,20 +32,20 @@ def run_model_and_get_history(model_path: str, url: str, steps: int) -> tuple[st
     source = Path(model_path).read_text()
 
     # Dynamically load the model module
-    spec = importlib.util.spec_from_file_location("generated_model", model_path)
+    spec = importlib.util.spec_from_file_location('generated_model', model_path)
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
 
     # Find the model class (first class with step_* methods)
     model_class = None
     for attr_name in dir(module):
         attr = getattr(module, attr_name)
-        if isinstance(attr, type) and any(m.startswith("step_") for m in dir(attr)):
+        if isinstance(attr, type) and any(m.startswith('step_') for m in dir(attr)):
             model_class = attr
             break
 
     if model_class is None:
-        print("Error: No PyOsmo model class found in the file")
+        print('Error: No PyOsmo model class found in the file')
         sys.exit(1)
 
     with sync_playwright() as p:
@@ -70,18 +70,18 @@ async def refine_model(model_path: str, url: str, steps: int) -> None:
     try:
         from claude_agent_sdk import Agent, AgentConfig
     except ImportError:
-        print("Error: claude-agent-sdk is required. Install with: pip install claude-agent-sdk")
+        print('Error: claude-agent-sdk is required. Install with: pip install claude-agent-sdk')
         sys.exit(1)
 
     from prompt_template import REFINEMENT_PROMPT
 
-    print(f"Running model {model_path}...")
+    print(f'Running model {model_path}...')
     source, history_json = run_model_and_get_history(model_path, url, steps)
 
-    print("Analyzing results and refining model...")
+    print('Analyzing results and refining model...')
 
     agent = Agent(
-        model="claude-sonnet-4-5-20250929",
+        model='claude-sonnet-4-5-20250929',
         config=AgentConfig(system_prompt=REFINEMENT_PROMPT),
     )
 
@@ -107,22 +107,22 @@ async def refine_model(model_path: str, url: str, steps: int) -> None:
 
     # Write refined model back
     Path(model_path).write_text(result)
-    print(f"Refined model written to {model_path}")
+    print(f'Refined model written to {model_path}')
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description="Refine a PyOsmo model based on test results")
-    parser.add_argument("model_file", help="Path to the PyOsmo model file")
-    parser.add_argument("--url", required=True, help="URL of the web application")
-    parser.add_argument("--steps", type=int, default=20, help="Steps per test (default: 20)")
-    parser.add_argument("--iterations", type=int, default=1, help="Refinement iterations (default: 1)")
+    parser = argparse.ArgumentParser(description='Refine a PyOsmo model based on test results')
+    parser.add_argument('model_file', help='Path to the PyOsmo model file')
+    parser.add_argument('--url', required=True, help='URL of the web application')
+    parser.add_argument('--steps', type=int, default=20, help='Steps per test (default: 20)')
+    parser.add_argument('--iterations', type=int, default=1, help='Refinement iterations (default: 1)')
     args = parser.parse_args()
 
     for i in range(args.iterations):
         if args.iterations > 1:
-            print(f"\n--- Refinement iteration {i + 1}/{args.iterations} ---")
+            print(f'\n--- Refinement iteration {i + 1}/{args.iterations} ---')
         asyncio.run(refine_model(args.model_file, args.url, args.steps))
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()

pyosmo/algorithm/balancing.py

@@ -8,14 +8,13 @@ class BalancingRandomAlgorithm(OsmoAlgorithm):
     In practise rare steps gets more priority when those are available"""
 
     def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
-        if self.random is None:
-            raise RuntimeError('Algorithm not initialized. Call initialize() first.')
+        rng = self._ensure_initialized()
         if len(choices) == 1:
             return choices[0]
         history_counts = [history.get_step_count(choice) for choice in choices]
         weights = [(sum(history_counts) - h) for h in history_counts]
         weights = [w - min(weights) + 1 for w in weights]
-        return self.random.choices(choices, weights=weights)[0]
+        return rng.choices(choices, weights=weights)[0]
 
 
 class BalancingAlgorithm(OsmoAlgorithm):

pyosmo/algorithm/base.py

@@ -22,6 +22,12 @@ def initialize(self, random: Random, model: OsmoModelCollector) -> None:
         self.random = random
         self.model = model
 
+    def _ensure_initialized(self) -> Random:
+        """Return the Random instance, raising if not yet initialized."""
+        if self.random is None:
+            raise RuntimeError('Algorithm not initialized. Call initialize() first.')
+        return self.random
+
     @abstractmethod
     def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
         raise Exception('This is just abstract class, not implementation')

pyosmo/algorithm/random.py

@@ -7,6 +7,4 @@ class RandomAlgorithm(OsmoAlgorithm):
     """Fully random algorithm"""
 
     def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
-        if self.random is None:
-            raise RuntimeError('Algorithm not initialized. Call initialize() first.')
-        return self.random.choice(choices)
+        return self._ensure_initialized().choice(choices)

pyosmo/algorithm/weighted.py

@@ -7,23 +7,20 @@ class WeightedAlgorithm(OsmoAlgorithm):
     """Weighted random algorithm"""
 
     def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
-        if self.random is None:
-            raise RuntimeError('Algorithm not initialized. Call initialize() first.')
-        return self.random.choices(choices, weights=[c.weight for c in choices])[0]
+        return self._ensure_initialized().choices(choices, weights=[c.weight for c in choices])[0]
 
 
 class WeightedBalancingAlgorithm(OsmoAlgorithm):
     """Weighted algorithm which balances based on history"""
 
     def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
-        if self.random is None:
-            raise RuntimeError('Algorithm not initialized. Call initialize() first.')
+        rng = self._ensure_initialized()
         weights = [c.weight for c in choices]
         normalized_weights = [float(i) / max(weights) for i in weights]
 
         history_counts = [history.get_step_count(choice) for choice in choices]
         if max(history_counts) == 0:
-            return self.random.choices(choices, weights=normalized_weights)[0]
+            return rng.choices(choices, weights=normalized_weights)[0]
 
         history_normalized_weights = [float(i) / max(history_counts) for i in history_counts]
 
@@ -36,4 +33,4 @@ def choose(self, history: OsmoHistory, choices: list[TestStep]) -> TestStep:
             temp_add = (abs(sum(total_weights)) + 0.2) / len(total_weights)
             total_weights = [temp_add + x for x in total_weights]
 
-        return self.random.choices(choices, weights=total_weights)[0]
+        return rng.choices(choices, weights=total_weights)[0]

pyosmo/config.py

@@ -25,66 +25,45 @@ class ConfigValidator:
     """Validates configuration values with comprehensive error messages."""
 
     @staticmethod
-    def validate_algorithm(algorithm: Any) -> None:
-        """Validate algorithm configuration.
-
-        Args:
-            algorithm: Algorithm to validate
-
-        Raises:
-            ConfigurationError: If algorithm is invalid
-        """
-        if algorithm is None:
-            raise ConfigurationError('Algorithm cannot be None. Please provide a valid OsmoAlgorithm instance.')
-
-        if not isinstance(algorithm, OsmoAlgorithm):
+    def _validate_type(value: Any, expected_type: type, name: str, available: str) -> None:
+        """Generic type validation for configuration values."""
+        if value is None:
             raise ConfigurationError(
-                f'Algorithm must be an instance of OsmoAlgorithm, '
-                f'got {type(algorithm).__name__}. '
-                f'Available algorithms: RandomAlgorithm, WeightedAlgorithm, BalancingAlgorithm.'
+                f'{name} cannot be None. Please provide a valid {expected_type.__name__} instance.'
             )
-
-    @staticmethod
-    def validate_end_condition(condition: Any, name: str = 'End condition') -> None:
-        """Validate end condition configuration.
-
-        Args:
-            condition: End condition to validate
-            name: Name of the condition for error messages
-
-        Raises:
-            ConfigurationError: If end condition is invalid
-        """
-        if condition is None:
-            raise ConfigurationError(f'{name} cannot be None. Please provide a valid OsmoEndCondition instance.')
-
-        if not isinstance(condition, OsmoEndCondition):
+        if not isinstance(value, expected_type):
             raise ConfigurationError(
-                f'{name} must be an instance of OsmoEndCondition, '
-                f'got {type(condition).__name__}. '
-                f'Available conditions: Length, Time, StepCoverage, Endless, And, Or.'
+                f'{name} must be an instance of {expected_type.__name__}, '
+                f'got {type(value).__name__}. '
+                f'Available: {available}.'
             )
 
     @staticmethod
-    def validate_error_strategy(strategy: Any, name: str = 'Error strategy') -> None:
-        """Validate error strategy configuration.
-
-        Args:
-            strategy: Error strategy to validate
-            name: Name of the strategy for error messages
+    def validate_algorithm(algorithm: Any) -> None:
+        ConfigValidator._validate_type(
+            algorithm,
+            OsmoAlgorithm,
+            'Algorithm',
+            'RandomAlgorithm, WeightedAlgorithm, BalancingAlgorithm',
+        )
 
-        Raises:
-            ConfigurationError: If error strategy is invalid
-        """
-        if strategy is None:
-            raise ConfigurationError(f'{name} cannot be None. Please provide a valid OsmoErrorStrategy instance.')
+    @staticmethod
+    def validate_end_condition(condition: Any, name: str = 'End condition') -> None:
+        ConfigValidator._validate_type(
+            condition,
+            OsmoEndCondition,
+            name,
+            'Length, Time, StepCoverage, Endless, And, Or',
+        )
 
-        if not isinstance(strategy, OsmoErrorStrategy):
-            raise ConfigurationError(
-                f'{name} must be an instance of OsmoErrorStrategy, '
-                f'got {type(strategy).__name__}. '
-                f'Available strategies: AlwaysRaise, AlwaysIgnore, IgnoreAsserts, AllowCount.'
-            )
+    @staticmethod
+    def validate_error_strategy(strategy: Any, name: str = 'Error strategy') -> None:
+        ConfigValidator._validate_type(
+            strategy,
+            OsmoErrorStrategy,
+            name,
+            'AlwaysRaise, AlwaysIgnore, IgnoreAsserts, AllowCount',
+        )
 
     @staticmethod
     def validate_seed(seed: Any) -> None:

pyosmo/error_strategy/always_ignore.py

@@ -10,6 +10,3 @@ class AlwaysIgnore(OsmoErrorStrategy):
 
     def failure_in_test(self, history: OsmoHistory, model: OsmoModelCollector, error: BaseException) -> None:
         pass
-
-    def failure_in_suite(self, history: OsmoHistory, model: OsmoModelCollector, error: BaseException) -> None:
-        pass

pyosmo/error_strategy/always_raise.py

@@ -10,6 +10,3 @@ class AlwaysRaise(OsmoErrorStrategy):
 
     def failure_in_test(self, history: OsmoHistory, model: OsmoModelCollector, error: BaseException) -> None:
         raise error
-
-    def failure_in_suite(self, history: OsmoHistory, model: OsmoModelCollector, error: BaseException) -> None:
-        raise error