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

waldiez/running/pre_run.py

@@ -7,12 +7,20 @@ import io
 import os
 import subprocess
 import sys
+import tempfile
+from pathlib import Path
 from typing import Callable
 
 from waldiez.models import Waldiez
 
 from .environment import in_virtualenv, is_root
-from .utils import strip_ansi
+from .utils import (
+    ensure_pip,
+    get_pip_install_location,
+    get_python_executable,
+    safe_filename,
+    strip_ansi,
+)
 
 
 class RequirementsMixin:
@@ -75,20 +83,12 @@ def install_requirements(
     """
     requirements_string = ", ".join(extra_requirements)
     printer(f"Installing requirements: {requirements_string}")
-    pip_install = [sys.executable, "-m", "pip", "install"]
-    break_system_packages = ""
-    if not in_virtualenv():  # it should  # pragma: no cover
-        # if not, let's try to install as user
-        # not sure if --break-system-packages is safe,
-        # but it might fail if we don't
-        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
-        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
-        if not is_root():
-            pip_install.append("--user")
-    if upgrade:  # pragma: no cover
-        pip_install.append("--upgrade")
-    pip_install.extend(extra_requirements)
-    # pylint: disable=too-many-try-statements
+    ensure_pip()
+    pip_install, break_system_packages, install_location = _before_pip(
+        extra_requirements, upgrade
+    )
+
+    # pylint: disable=too-many-try-statements,broad-exception-caught
     try:
         with subprocess.Popen(
             pip_install,
@@ -97,17 +97,26 @@ def install_requirements(
         ) as proc:
             if proc.stdout:  # pragma: no branch
                 for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
-                    printer(strip_ansi(line.strip()))
+                    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"):
-                    printer(strip_ansi(line.strip()))
+                    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:
+                printer(
+                    f"Package installation failed with exit code {return_code}"
+                )
+
+    except Exception as e:
+        printer(f"Failed to install requirements: {e}")
     finally:
-        if not in_virtualenv():  # pragma: no cover
-            # restore the old env var
-            if break_system_packages:
-                os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
-            else:
-                del os.environ["PIP_BREAK_SYSTEM_PACKAGES"]
+        _after_pip(break_system_packages, install_location)
 
 
 async def a_install_requirements(
@@ -126,34 +135,147 @@ async def a_install_requirements(
     printer : Callable[..., None]
         The printer function to use, defaults to print.
     """
+    pip_install, break_system_packages, install_location = _before_pip(
+        extra_requirements, upgrade=upgrade
+    )
     requirements_string = ", ".join(extra_requirements)
     printer(f"Installing requirements: {requirements_string}")
-    pip_install = [sys.executable, "-m", "pip", "install"]
-    break_system_packages = ""
-    if not in_virtualenv():  # pragma: no cover
-        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
-        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
-        if not is_root():
-            pip_install.extend(["--user"])
-    if upgrade:  # pragma: no cover
-        pip_install.append("--upgrade")
-    pip_install.extend(extra_requirements)
-    # pylint: disable=too-many-try-statements
+    # pylint: disable=too-many-try-statements,broad-exception-caught
     try:
         proc = await asyncio.create_subprocess_exec(
             *pip_install,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
         )
-        if proc.stdout:
-            async for line in proc.stdout:
-                printer(strip_ansi(line.decode().strip()))
-        if proc.stderr:
-            async for line in proc.stderr:
-                printer(strip_ansi(line.decode().strip()))
+
+        async def _pump_stream(stream: asyncio.StreamReader | None) -> None:
+            if not stream:
+                return
+            async for raw in stream:
+                text = strip_ansi(raw.decode().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:
+            printer(
+                f"Package installation failed with exit code {proc.returncode}"
+            )
+
+    except Exception as e:
+        printer(f"Failed to install requirements: {e}")
     finally:
-        if not in_virtualenv():  # pragma: no cover
-            if break_system_packages:
-                os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
-            else:
-                del os.environ["PIP_BREAK_SYSTEM_PACKAGES"]
+        _after_pip(break_system_packages, install_location)
+
+
+def _before_pip(
+    packages: set[str],
+    upgrade: bool,
+) -> tuple[list[str], str, str | None]:
+    """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, str | None]
+        The pip command, break_system_packages flag, and install location.
+    """
+    ensure_pip()
+    pip_install = [
+        get_python_executable(),
+        "-m",
+        "pip",
+        "install",
+        "--disable-pip-version-check",
+        "--no-input",
+    ]
+    install_location = get_pip_install_location()
+    break_system_packages = ""
+
+    if install_location:
+        pip_install += ["--target", install_location]
+    elif not in_virtualenv():  # it should  # pragma: no cover
+        # if not, let's try to install as user
+        # not sure if --break-system-packages is safe,
+        # but it might fail if we don't
+        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
+        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
+        if not is_root():
+            pip_install.append("--user")
+
+    if upgrade:  # pragma: no cover
+        pip_install.append("--upgrade")
+
+    pip_install.extend(packages)
+    return pip_install, break_system_packages, install_location
+
+
+def _after_pip(
+    break_system_packages: str, install_location: str | None
+) -> None:
+    """Restore environment variables after pip installation.
+
+    Parameters
+    ----------
+    break_system_packages : str
+        The original value of PIP_BREAK_SYSTEM_PACKAGES.
+    install_location : str | None
+        The install location used (None if default).
+    """
+    if install_location is None and not 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)
+
+
+def dump_waldiez(waldiez: Waldiez, output_path: str | Path | None) -> Path:
+    """Dump waldiez flow to a file.
+
+    Parameters
+    ----------
+    waldiez : Waldiez
+        The waldiez instance to dump.
+    output_path : str | Path | None
+        Optional output path to determine the directory to save the flow to.
+
+    Returns
+    -------
+    Path
+        The path to the generated file.
+    """
+    file_path = Path(output_path) if output_path else None
+    if file_path:
+        file_name = file_path.name
+        if not file_name.endswith(".waldiez"):
+            file_path.with_suffix(".waldiez")
+
+    else:
+        full_name = waldiez.name
+        file_name = safe_filename(full_name, "waldiez")
+    file_dir: Path
+    if file_path:
+        file_dir = file_path if file_path.is_dir() else file_path.parent
+    else:
+        file_dir = Path(tempfile.mkdtemp())
+    file_dir.mkdir(parents=True, exist_ok=True)
+    output_path = file_dir / file_name
+    with output_path.open(
+        "w", encoding="utf-8", errors="replace", newline="\n"
+    ) as f_open:
+        f_open.write(waldiez.model_dump_json())
+    return output_path

waldiez/running/standard_runner.py

@@ -2,7 +2,7 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
 # pyright: reportUnknownMemberType=false, reportAttributeAccessIssue=false
-# pylint: disable=duplicate-code
+# pylint: disable=duplicate-code,too-few-public-methods
 """Run a waldiez flow.
 
 The flow is first converted to an autogen flow with agents, chats, and tools.
@@ -59,9 +59,8 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
         self._event_count = 0
         self._processed_events = 0
 
-    # pylint: disable=no-self-use
-    # noinspection PyMethodMayBeStatic
-    def print(self, *args: Any, **kwargs: Any) -> None:
+    @staticmethod
+    def print(*args: Any, **kwargs: Any) -> None:
         """Print.
 
         Parameters
@@ -71,7 +70,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
         **kwargs : Any
             Keyword arguments to print.
         """
-        WaldiezBaseRunner._print(*args, **kwargs)
+        WaldiezBaseRunner.print(*args, **kwargs)
 
     # pylint: disable=unused-argument
     def _run(
@@ -214,6 +213,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
                         "workflow event processing (async)"
                     )
                     return []
+                # noinspection DuplicatedCode
                 if self.structured_io:
                     stream = StructuredIOStream(
                         uploads_root=uploads_root, is_async=True