waldiez 0.5.10__py3-none-any.whl → 0.6.1__py3-none-any.whl

waldiez/utils/python_manager.py

@@ -0,0 +1,415 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-many-try-statements,broad-exception-caught
+# pylint: disable=import-error
+"""Python manager class."""
+
+import asyncio
+import io
+import json
+import os
+import platform
+import re
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+WALDIEZ_SITE_PACKAGES = "WALDIEZ_SITE_PACKAGES"
+WALDIEZ_APP_ROOT = "WALDIEZ_APP_ROOT"
+
+
+# noinspection PyBroadException
+class PythonManager:
+    """Python manager."""
+
+    _app_dir: Path | None
+
+    def __init__(self) -> None:
+        self.system = platform.system().lower()
+        self.is_frozen = bool(getattr(sys, "frozen", False))
+        self._app_dir = None
+
+    @property
+    def app_dir(self) -> Path:
+        """Directory to read packaged resources from.
+
+        Returns
+        -------
+        Path
+            The path to the application installation directory.
+        """
+        return self._get_app_dir()
+
+    @property
+    def site_packages_directory(self) -> Path | None:
+        """Determine the best location to install packages.
+
+        Returns
+        -------
+        str | None
+            The installation target directory, or None for default.
+        """
+        return self._get_site_packages_path()
+
+    def get_python_executable(self) -> str:
+        """Get the appropriate Python executable path.
+
+        Returns
+        -------
+        str
+            The path to the appropriate Python executable path.
+        """
+        if not self.is_frozen:
+            return sys.executable
+        # Check for bundled Python in installation directory
+        app_dir = self.app_dir
+        if self.system == "windows":
+            candidates = [
+                app_dir / "bundled_python" / "Scripts" / "python.exe",
+                app_dir / "bundled_python" / "Scripts" / "python3.exe",
+                app_dir / "bundled_python" / "python.exe",
+            ]
+        else:
+            candidates = [
+                app_dir / "bundled_python" / "bin" / "python3",
+                app_dir / "bundled_python" / "bin" / "python",
+                app_dir / "bundled_python" / "python3",
+            ]
+
+        for candidate in candidates:
+            if candidate.exists():
+                return str(candidate)
+        # Fallback to system Python
+        return sys.executable
+
+    def list_installed_packages(self) -> list[dict[str, str]]:
+        """List packages in our managed environment.
+
+        Returns
+        -------
+        list[dict[str,str]]
+            The locally installed packages.
+        """
+        python_exe = self.get_python_executable()
+
+        env = os.environ.copy()
+        if self.site_packages_directory:
+            env["PYTHONPATH"] = str(self.site_packages_directory)
+
+        cmd = [python_exe, "-m", "pip", "list", "--format=json"]
+        try:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                env=env,
+                check=True,
+            )
+            if result.returncode == 0:
+                return json.loads(result.stdout)
+        except Exception:
+            pass
+
+        return []
+
+    def get_debug_info(self) -> dict[str, Any]:
+        """Get comprehensive debug information.
+
+        Returns
+        -------
+        dict[str, Any]
+            Debug info about the paths and the packages.
+        """
+        return {
+            "system": self.system,
+            "is_frozen": self.is_frozen,
+            "site_packages_directory": (
+                str(self.site_packages_directory)
+                if self.site_packages_directory
+                else None
+            ),
+            "python_executable": self.get_python_executable(),
+            "python_version": sys.version,
+            "in_virtualenv": self.in_virtualenv(),
+            "sys_path": sys.path,
+            "pythonpath": os.environ.get("PYTHONPATH", ""),
+        }
+
+    def pip_install(
+        self,
+        packages: set[str],
+        upgrade: bool = False,
+        printer: Callable[..., None] = print,
+    ) -> None:
+        """Install packages.
+
+        Parameters
+        ----------
+        packages : set[str]
+            The packages to install.
+        upgrade : bool
+            Upgrade existing or not.
+        printer : Callable[..., None]
+            The callable to use for printing the process' output
+        """
+        pip_install_cmd, break_system_packages = self._before_pip(
+            packages, upgrade
+        )
+        try:
+            with subprocess.Popen(
+                pip_install_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            ) as proc:
+                if proc.stdout:  # pragma: no branch
+                    for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
+                        stripped_line = strip_ansi(line.strip())
+                        if stripped_line:  # Only print non-empty lines
+                            printer(stripped_line)
+                if proc.stderr:  # pragma: no branch
+                    for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
+                        stripped_line = strip_ansi(line.strip())
+                        if stripped_line:  # Only print non-empty lines
+                            printer(stripped_line)
+
+                # Wait for process to complete and check return code
+                return_code = proc.wait()
+                if return_code != 0:  # pragma: no cover
+                    msg = (
+                        "Package installation failed "
+                        f"with exit code {return_code}"
+                    )
+                    printer(msg)
+
+        except Exception as e:  # pragma: no cover
+            printer(f"Failed to install requirements: {e}")
+        finally:
+            self._after_pip(break_system_packages)
+
+    async def a_pip_install(
+        self,
+        packages: set[str],
+        upgrade: bool = False,
+        printer: Callable[..., None] = print,
+    ) -> None:
+        """Install packages asynchronously.
+
+        Parameters
+        ----------
+        packages : set[str]
+            The packages to install.
+        upgrade : bool
+            Upgrade existing or not.
+        printer : Callable[..., None]
+            The callable to use for printing the process' output
+        """
+        pip_install, break_system_packages = self._before_pip(
+            packages, upgrade=upgrade
+        )
+        requirements_string = ", ".join(packages)
+        printer(f"Installing requirements: {requirements_string}")
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *pip_install,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+
+            async def _pump_stream(stream: asyncio.StreamReader | None) -> None:
+                if not stream:  # pragma: no cover
+                    return
+                async for raw in stream:
+                    text = strip_ansi(raw.decode(errors="replace").rstrip())
+                    if text:
+                        printer(text)
+
+            # Create tasks for concurrent execution
+            tasks: list[asyncio.Task[int | None]] = [
+                asyncio.create_task(_pump_stream(proc.stdout)),
+                asyncio.create_task(_pump_stream(proc.stderr)),
+                asyncio.create_task(proc.wait()),
+            ]
+            await asyncio.gather(*tasks, return_exceptions=True)
+            if proc.returncode != 0:  # pragma: no cover
+                msg = (
+                    "Package installation failed "
+                    f"with exit code {proc.returncode}"
+                )
+                printer(msg)
+
+        except Exception as e:
+            printer(f"Failed to install requirements: {e}")
+        finally:
+            self._after_pip(break_system_packages)
+
+    # noinspection TryExceptPass
+    @staticmethod
+    def _ensure_pip() -> None:  # pragma: no cover
+        """Make sure `python -m pip` works (bootstrap if needed)."""
+        # pylint: disable=import-outside-toplevel,unused-import
+        try:
+            import pip  # noqa: F401  # pyright: ignore[reportUnusedImport]
+
+            return
+        except Exception:
+            pass
+        try:
+            import ensurepip
+
+            ensurepip.bootstrap(upgrade=True)
+        except Exception:
+            # If bootstrap fails,
+            # we'll still attempt `-m pip` and surface errors.
+            pass
+
+    def _get_app_dir(self) -> Path:
+        if self._app_dir is None:
+            from_env_str = os.environ.get(WALDIEZ_APP_ROOT, "")
+            if from_env_str:
+                from_env_path = Path(from_env_str).resolve()
+                if from_env_path.is_dir():
+                    self._app_dir = from_env_path
+                    return self._app_dir
+            if self.is_frozen and hasattr(sys, "_MEIPASS"):  # PyInstaller
+                self._app_dir = Path(
+                    getattr(sys, "_MEIPASS", Path(sys.executable).parent)
+                )
+                return self._app_dir
+        # dev: package root
+        self._app_dir = Path(__file__).parent.parent
+        return self._app_dir
+
+    def _get_site_packages_path(self) -> Path | None:
+        from_env_str = os.environ.get(WALDIEZ_SITE_PACKAGES, "")
+        if from_env_str:
+            from_env_path = Path(from_env_str).resolve()
+            if from_env_path.is_dir():
+                return from_env_path
+        if self.is_frozen:
+            # Use the bundled site-packages if available
+            bundled_sp = self.app_dir / "bundled_python" / "site-packages"
+            if bundled_sp.exists():
+                return bundled_sp
+        return None
+
+    def _before_pip(
+        self,
+        packages: set[str],
+        upgrade: bool,
+    ) -> tuple[list[str], str]:
+        """Gather the pip command for installing requirements.
+
+        Parameters
+        ----------
+        packages : set[str]
+            The packages to install.
+        upgrade : bool
+            Whether to upgrade the packages.
+
+        Returns
+        -------
+        tuple[list[str], str]
+            The pip command, and the break_system_packages flag.
+        """
+        self._ensure_pip()
+        pip_install = [
+            self.get_python_executable(),
+            "-m",
+            "pip",
+            "install",
+            "--disable-pip-version-check",
+            "--no-input",
+        ]
+        install_location = self.site_packages_directory
+        break_system_packages = ""
+
+        if install_location:
+            pip_install += ["--target", str(install_location)]
+        elif not self.in_virtualenv():  # pragma: no cover
+            # it should, if not, let's try to install as user
+            if not is_root():
+                pip_install.append("--user")
+            break_system_packages = os.environ.get(
+                "PIP_BREAK_SYSTEM_PACKAGES", ""
+            )
+            os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
+
+        if upgrade:  # pragma: no cover
+            pip_install.append("--upgrade")
+
+        pip_install.extend(sorted(packages))
+        return pip_install, break_system_packages
+
+    def _after_pip(
+        self,
+        break_system_packages: str,
+    ) -> None:
+        """Restore environment variables after pip installation.
+
+        Parameters
+        ----------
+        break_system_packages : str
+            The original value of PIP_BREAK_SYSTEM_PACKAGES.
+        """
+        if (
+            not self.site_packages_directory and not self.in_virtualenv()
+        ):  # pragma: no cover
+            # restore the old env var
+            if break_system_packages:
+                os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
+            else:
+                # Use pop to avoid KeyError if the key doesn't exist
+                os.environ.pop("PIP_BREAK_SYSTEM_PACKAGES", None)
+
+    @staticmethod
+    def in_virtualenv() -> bool:
+        """Check if we are inside a virtualenv.
+
+        Returns
+        -------
+        bool
+            True if inside a virtualenv, False otherwise.
+        """
+        return hasattr(sys, "real_prefix") or (
+            hasattr(sys, "base_prefix")
+            and os.path.realpath(sys.base_prefix)
+            != os.path.realpath(sys.prefix)
+        )
+
+
+def is_root() -> bool:  # pragma: no cover  # os specific
+    """Check if the script is running as root/administrator.
+
+    Returns
+    -------
+    bool
+        True if running as root/administrator, False otherwise.
+    """
+    # pylint: disable=import-outside-toplevel,line-too-long,no-member
+    if os.name == "nt":
+        try:
+            import ctypes
+
+            return ctypes.windll.shell32.IsUserAnAdmin() != 0  # type: ignore[unused-ignore,attr-defined]  # noqa: E501
+        except Exception:
+            return False
+    else:
+        return os.getuid() == 0
+
+
+def strip_ansi(text: str) -> str:
+    """Remove ANSI escape sequences from text.
+
+    Parameters
+    ----------
+    text : str
+        The text to strip.
+
+    Returns
+    -------
+    str
+        The text without ANSI escape sequences.
+    """
+    ansi_pattern = re.compile(r"\x1b\[[0-9;]*m|\x1b\[.*?[@-~]")
+    return ansi_pattern.sub("", text)

waldiez/ws/__init__.py

@@ -18,7 +18,7 @@ from .errors import (
     UnsupportedActionError,
     WaldiezServerError,
 )
-from .server import WaldiezWsServer, run_server
+from .server import HAS_WEBSOCKETS, WaldiezWsServer, run_server
 from .session_manager import SessionManager
 from .utils import (
     ConnectionManager,
@@ -38,13 +38,14 @@ def add_ws_app(app: typer.Typer) -> None:
     app : typer.Typer
         The Typer application instance.
     """
-    app.registered_commands.append(
-        CommandInfo(
-            name="ws",
-            help="Start the Waldiez WebSocket server.",
-            callback=serve,
+    if HAS_WEBSOCKETS:
+        app.registered_commands.append(
+            CommandInfo(
+                name="ws",
+                help="Start the Waldiez WebSocket server.",
+                callback=serve,
+            )
         )
-    )
 
 
 __all__ = [

waldiez/ws/_file_handler.py

@@ -3,8 +3,6 @@
 
 """Files related request handler."""
 
-from __future__ import annotations
-
 import json
 import logging
 from pathlib import Path
@@ -25,7 +23,7 @@ class FileRequestHandler:
     """Handles file-related requests."""
 
     @staticmethod
-    def handle_save_flow_request(
+    def handle_save_request(
         msg: SaveFlowRequest,
         workspace_dir: Path,
         client_id: str,
@@ -49,10 +47,10 @@ class FileRequestHandler:
         dict[str, Any]
             The response dictionary.
         """
-        filename = msg.filename or f"waldiez_{client_id}.waldiez"
+        path = msg.path or f"waldiez_{client_id}.waldiez"
         try:
             output_path = resolve_output_path(
-                filename,
+                path,
                 workspace_dir=workspace_dir,
                 expected_ext="waldiez",
             )
@@ -60,28 +58,28 @@ class FileRequestHandler:
             logger.error("Error resolving output path: %s", exc)
             return SaveFlowResponse.fail(
                 error=f"Invalid output path: {exc}",
-                file_path=filename,
+                path=path,
             ).model_dump(mode="json")
         # pylint: disable=too-many-try-statements
         try:
-            if output_path.exists() and not msg.force_overwrite:
+            if output_path.exists() and not msg.force:
                 return SaveFlowResponse.fail(
                     error=f"File exists: {output_path}",
                     file_path=str(output_path.relative_to(workspace_dir)),
                 ).model_dump(mode="json")
 
             # Parent dir already created by resolve_output_path
-            output_path.write_text(msg.flow_data, encoding="utf-8")
+            output_path.write_text(msg.data, encoding="utf-8", newline="\n")
 
             return SaveFlowResponse.ok(
-                file_path=str(output_path.relative_to(workspace_dir))
+                path=str(output_path.relative_to(workspace_dir))
             ).model_dump(mode="json")
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.error("Error saving flow: %s", e)
             return SaveFlowResponse.fail(error=str(e)).model_dump(mode="json")
 
     @staticmethod
-    def handle_convert_workflow_request(
+    def handle_convert_request(
         msg: ConvertWorkflowRequest,
         client_id: str,
         workspace_dir: Path,
@@ -105,26 +103,26 @@ class FileRequestHandler:
         dict[str, Any]
             The response dictionary.
         """
-        target_format = (msg.target_format or "").strip().lower()
+        target_format = (msg.format or "").strip().lower()
         if target_format not in {"py", "ipynb"}:
             return ConvertWorkflowResponse.fail(
                 error=f"Unsupported target format: {target_format}",
-                target_format=target_format,
+                format=target_format,
             ).model_dump(mode="json")
 
         try:
-            waldiez_data = Waldiez.from_dict(json.loads(msg.flow_data))
+            waldiez_data = Waldiez.from_dict(json.loads(msg.data))
         except Exception as e:  # pylint: disable=broad-exception-caught
             return ConvertWorkflowResponse.fail(
                 error=f"Invalid flow_data: {e}",
-                target_format=target_format,
+                format=target_format,
             ).model_dump(mode="json")
 
         try:
             # Use normalized target_format for default name
-            filename = msg.output_path or f"waldiez_{client_id}.{target_format}"
+            path = msg.path or f"waldiez_{client_id}.{target_format}"
             output_path = resolve_output_path(
-                filename,
+                path,
                 workspace_dir=workspace_dir,
                 expected_ext=target_format,
             )
@@ -132,7 +130,7 @@ class FileRequestHandler:
             logger.error("Error resolving output path: %s", exc)
             return ConvertWorkflowResponse.fail(
                 error=f"Invalid output path: {exc}",
-                target_format=target_format,
+                format=target_format,
             ).model_dump(mode="json")
 
         try:
@@ -140,13 +138,13 @@ class FileRequestHandler:
             exporter.export(path=output_path, force=True, structured_io=True)
 
             return ConvertWorkflowResponse.ok(
-                target_format=target_format,
-                output_path=str(output_path.relative_to(workspace_dir)),
+                format=target_format,
+                path=str(output_path.relative_to(workspace_dir)),
             ).model_dump(mode="json")
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.error("Error converting workflow: %s", e)
             return ConvertWorkflowResponse.fail(
-                error=str(e), target_format=target_format
+                error=str(e), format=target_format
             ).model_dump(mode="json")
 
 

waldiez/ws/_mock.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+# pylint: disable=invalid-name,line-too-long,unused-argument,too-few-public-methods,no-self-use
+# pylint: disable=missing-class-docstring,missing-function-docstring,missing-return-doc
+# flake8: noqa: E501, D101, D102, D106
+# pyright:  reportUnusedParameter=false, reportUninitializedInstanceVariable=false
+"""Mock websockets for linters."""
+
+from typing import Any  # pragma: no cover
+
+
+# noinspection PyPep8Naming, PyMethodMayBeStatic,PyUnusedLocal
+class websockets:  # pragma: no cover
+    # noinspection PyMethodMayBeStatic
+    class ClientConnection:
+        async def __aenter__(self) -> "websockets.ClientConnection":
+            return self
+
+        async def send(self, data: Any) -> None:
+            pass
+
+        async def close(self) -> None:
+            pass
+
+        async def recv(self) -> str:
+            return ""
+
+    # noinspection PyPep8Naming
+    class exceptions:
+        class ConnectionClosedError(Exception): ...
+
+    class ConnectionClosedOK(Exception): ...
+
+    class Server:
+        def close(self) -> None:
+            pass
+
+        async def wait_closed(self) -> None:
+            pass
+
+    class ServerConnection:
+        remote_address: str
+
+        # noinspection PyPep8Naming
+        class request:
+            headers: dict[str, str]
+
+        async def close(
+            self, code: int = 1000, reason: str = "Normal Closure"
+        ) -> None:
+            pass
+
+        # noinspection PyPep8Naming
+        async def ConnectionClosedOK(self, *args: Any, **kwargs: Any) -> None:
+            pass
+
+        async def send(self, data: Any) -> None:
+            pass
+
+    @staticmethod
+    async def serve(*args: Any, **kwargs: Any) -> "websockets.Server":
+        return websockets.Server()
+
+    @staticmethod
+    async def connect(
+        *args: Any, **kwargs: Any
+    ) -> "websockets.ClientConnection":
+        return websockets.ClientConnection()
+
+    class ConnectionClosed(Exception):
+        pass
+
+    class WebSocketException(Exception):
+        pass