Refactor into standalone function

Author: fperez

packages/jupyter-ai/jupyter_ai/personas/persona_manager.py

@@ -145,8 +145,8 @@ def _init_persona_classes(self) -> None:
                 + "Please verify your server configuration and open a new issue on our GitHub repo if this warning persists."
             )
 
-        # TODO: check whether in the end we need a full class for this, or if we can just use a simple function that returns a list of persona classes from the local filesystem.
-        persona_classes.extend(LocalPersonaLoader(self.root_dir).load_persona_classes())
+        # Load persona classes from local filesystem
+        persona_classes.extend(load_persona_classes_from_directory(self.root_dir))
 
         PersonaManager._persona_classes = persona_classes
 
@@ -297,114 +297,66 @@ def get_mcp_config(self) -> dict[str, Any]:
             return self._mcp_config_loader.get_config(jdir)
 
 
-class LocalPersonaLoader(LoggingConfigurable):
+def load_persona_classes_from_directory(root_dir: str) -> list[type[BasePersona]]:
     """
-    Load _persona class declarations_ from the local filesystem.
+    Load _persona class declarations_ from Python files in the local filesystem.
 
-    Those class declarations are then used to instantiate personas by the `PersonaManager`.
+    Those class declarations are then used to instantiate personas by the
+    `PersonaManager`.
 
-    TODO: wire to local .jupyter directory system, for now just use the current
-    directory.
-    """
-
-    def __init__(
-        self,
-        *args,
-        root_dir: str | None = None,
-        **kwargs,
-    ):
-        # Forward other arguments to parent class
-        super().__init__(*args, **kwargs)
+    Scans the root_dir for .py files containing "persona" in their name,
+    dynamically imports them, and extracts any class declarations that are
+    subclasses of BasePersona.
 
-        # Set default root directory to current working directory if not provided
-        self.root_dir = root_dir or os.getcwd()
+    Args:
+        root_dir: Directory to scan for persona Python files
 
-        self.log.info(
-            f"LocalPersonaLoader initialized with root directory: {self.root_dir}"
-        )
-
-    def load_persona_classes(self) -> list[type[BasePersona]]:
-        """
-        Loads persona classes from Python files in the local filesystem.
-
-        Scans the root_dir for .py files, dynamically imports them, and extracts
-        any class declarations that are subclasses of BasePersona.
-        """
-        persona_classes: list[type[BasePersona]] = []
-
-        # Check if root directory exists
-        if not os.path.exists(self.root_dir):
-            self.log.info(f"Root directory does not exist: {self.root_dir}")
-            return persona_classes
+    Returns:
+        List of BasePersona subclasses found in the directory
+    """
+    persona_classes: list[type[BasePersona]] = []
 
-        # Find all .py files in the root directory that contain "persona" in the name
-        all_py_files = glob(os.path.join(self.root_dir, "*.py"))
-        py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()]
+    # Check if root directory exists
+    if not os.path.exists(root_dir):
+        return persona_classes
 
-        if not py_files:
-            self.log.info(
-                f"No Python files with 'persona' in the name found in directory: {self.root_dir}"
-            )
-            return persona_classes
+    # Find all .py files in the root directory that contain "persona" in the name
+    all_py_files = glob(os.path.join(root_dir, "*.py"))
+    py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()]
 
-        self.log.info(
-            f"Found {len(py_files)} Python files with 'persona' in the name in {self.root_dir}"
-        )
-        self.log.info("PENDING: Loading persona classes from local Python files...")
-        start_time_ns = time_ns()
+    if not py_files:
+        return persona_classes
 
-        for py_file in py_files:
-            try:
-                # Get module name from file path
-                module_name = Path(py_file).stem
-
-                # Skip if module name starts with underscore (private modules)
-                if module_name.startswith("_"):
-                    continue
-
-                # Create module spec and load the module
-                spec = importlib.util.spec_from_file_location(module_name, py_file)
-                if spec is None or spec.loader is None:
-                    self.log.warning(f"  - Unable to create module spec for {py_file}")
-                    continue
-
-                module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(module)
-
-                # Find all classes in the module that are BasePersona subclasses
-                module_persona_classes = []
-                for name, obj in inspect.getmembers(module, inspect.isclass):
-                    # Check if it's a subclass of BasePersona but not BasePersona itself
-                    if (
-                        issubclass(obj, BasePersona)
-                        and obj is not BasePersona
-                        and obj.__module__ == module_name
-                    ):
-                        module_persona_classes.append(obj)
-
-                if module_persona_classes:
-                    persona_classes.extend(module_persona_classes)
-                    class_names = [cls.__name__ for cls in module_persona_classes]
-                    self.log.info(
-                        f"  - Loaded {len(module_persona_classes)} persona class(es) from '{py_file}': {class_names}"
-                    )
-                else:
-                    self.log.debug(f"  - No persona classes found in '{py_file}'")
+    for py_file in py_files:
+        try:
+            # Get module name from file path
+            module_name = Path(py_file).stem
 
-            except Exception:
-                # On exception, log an error and continue
-                # This mirrors the error handling pattern from entry point loading
-                self.log.exception(
-                    f"  - Unable to load persona classes from '{py_file}' due to an exception printed below."
-                )
+            # Skip if module name starts with underscore (private modules)
+            if module_name.startswith("_"):
                 continue
 
-        if len(persona_classes) > 0:
-            elapsed_time_ms = (time_ns() - start_time_ns) // 1_000_000
-            self.log.info(
-                f"SUCCESS: Loaded {len(persona_classes)} persona classes from local filesystem. Time elapsed: {elapsed_time_ms}ms."
-            )
-        else:
-            self.log.info("No persona classes found in local filesystem.")
+            # Create module spec and load the module
+            spec = importlib.util.spec_from_file_location(module_name, py_file)
+            if spec is None or spec.loader is None:
+                continue
 
-        return persona_classes
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+
+            # Find all classes in the module that are BasePersona subclasses
+            for name, obj in inspect.getmembers(module, inspect.isclass):
+                # Check if it's a subclass of BasePersona but not BasePersona itself
+                if (
+                    issubclass(obj, BasePersona)
+                    and obj is not BasePersona
+                    and obj.__module__ == module_name
+                ):
+                    persona_classes.append(obj)
+
+        except Exception:
+            # On exception, continue to next file
+            # This mirrors the error handling pattern from entry point loading
+            continue
+
+    return persona_classes

packages/jupyter-ai/jupyter_ai/tests/test_personas.py

@@ -7,7 +7,7 @@
 
 import pytest
 from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults
-from jupyter_ai.personas.persona_manager import LocalPersonaLoader
+from jupyter_ai.personas.persona_manager import load_persona_classes_from_directory
 
 
 @pytest.fixture
@@ -17,13 +17,12 @@ def tmp_persona_dir():
         yield Path(temp_dir)
 
 
-class TestLocalPersonaLoader:
-    """Test cases for LocalPersonaLoader class."""
+class TestLoadPersonaClassesFromDirectory:
+    """Test cases for load_persona_classes_from_directory function."""
 
     def test_empty_directory_returns_empty_list(self, tmp_persona_dir):
         """Test that an empty directory returns an empty list of persona classes."""
-        loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir))
-        result = loader.load_persona_classes()
+        result = load_persona_classes_from_directory(str(tmp_persona_dir))
         assert result == []
 
     def test_non_persona_file_returns_empty_list(self, tmp_persona_dir):
@@ -32,8 +31,7 @@ def test_non_persona_file_returns_empty_list(self, tmp_persona_dir):
         non_persona_file = tmp_persona_dir / "no_personas.py"
         non_persona_file.write_text("pass")
 
-        loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir))
-        result = loader.load_persona_classes()
+        result = load_persona_classes_from_directory(str(tmp_persona_dir))
         assert result == []
 
     def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir):
@@ -53,28 +51,18 @@ def process_message(self, message):
 """
         persona_file.write_text(persona_content)
 
-        loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir))
-        result = loader.load_persona_classes()
+        result = load_persona_classes_from_directory(str(tmp_persona_dir))
 
         assert len(result) == 1
         assert result[0].__name__ == "TestPersona"
         assert issubclass(result[0], BasePersona)
 
-    def test_bad_persona_file_logs_exception_returns_empty_list(
-        self, tmp_persona_dir, caplog
-    ):
-        """Test that a file with syntax errors logs an exception and returns empty list."""
+    def test_bad_persona_file_returns_empty_list(self, tmp_persona_dir):
+        """Test that a file with syntax errors returns empty list."""
         # Create a file with invalid Python code
         bad_persona_file = tmp_persona_dir / "bad_persona.py"
         bad_persona_file.write_text("1/0")
 
-        loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir))
-        result = loader.load_persona_classes()
+        result = load_persona_classes_from_directory(str(tmp_persona_dir))
 
         assert result == []
-        # Check that an error was logged
-        assert any(
-            "Unable to load persona classes from" in record.message
-            for record in caplog.records
-            if record.levelname in ["ERROR", "EXCEPTION"]
-        )