Added python funcs to computer

libs/computer/computer/computer.py

@@ -722,3 +722,173 @@ class Computer:
             tuple[float, float]: (x, y) coordinates in screenshot space
         """
         return await self.interface.to_screenshot_coordinates(x, y)
+
+
+    # Add virtual environment management functions to computer interface
+    async def venv_install(self, venv_name: str, requirements: list[str]) -> tuple[str, str]:
+        """Install packages in a virtual environment.
+        
+        Args:
+            venv_name: Name of the virtual environment
+            requirements: List of package requirements to install
+            
+        Returns:
+            Tuple of (stdout, stderr) from the installation command
+        """
+        requirements = requirements or []
+
+        # Create virtual environment if it doesn't exist
+        venv_path = f"~/.venvs/{venv_name}"
+        create_cmd = f"mkdir -p ~/.venvs && python3 -m venv {venv_path}"
+        
+        # Check if venv exists, if not create it
+        check_cmd = f"test -d {venv_path} || ({create_cmd})"
+        _, _ = await self.interface.run_command(check_cmd)
+        
+        # Install packages
+        requirements_str = " ".join(requirements)
+        install_cmd = f". {venv_path}/bin/activate && pip install {requirements_str}"
+        return await self.interface.run_command(install_cmd)
+    
+    async def venv_cmd(self, venv_name: str, command: str) -> tuple[str, str]:
+        """Execute a shell command in a virtual environment.
+        
+        Args:
+            venv_name: Name of the virtual environment
+            command: Shell command to execute in the virtual environment
+            
+        Returns:
+            Tuple of (stdout, stderr) from the command execution
+        """
+        venv_path = f"~/.venvs/{venv_name}"
+        
+        # Check if virtual environment exists
+        check_cmd = f"test -d {venv_path}"
+        stdout, stderr = await self.interface.run_command(check_cmd)
+        
+        if stderr or "test:" in stdout:  # venv doesn't exist
+            return "", f"Virtual environment '{venv_name}' does not exist. Create it first using venv_install."
+        
+        # Activate virtual environment and run command
+        full_command = f". {venv_path}/bin/activate && {command}"
+        return await self.interface.run_command(full_command)
+    
+    async def venv_exec(self, venv_name: str, python_func, *args, **kwargs):
+        """Execute Python function in a virtual environment using source code extraction.
+        
+        Args:
+            venv_name: Name of the virtual environment
+            python_func: A callable function to execute
+            *args: Positional arguments to pass to the function
+            **kwargs: Keyword arguments to pass to the function
+            
+        Returns:
+            The result of the function execution, or raises any exception that occurred
+        """
+        import base64
+        import inspect
+        import json
+        import textwrap
+        
+        try:
+            # Get function source code using inspect.getsource
+            source = inspect.getsource(python_func)
+            # Remove common leading whitespace (dedent)
+            func_source = textwrap.dedent(source).strip()
+            
+            # Get function name for execution
+            func_name = python_func.__name__
+            
+            # Serialize args and kwargs as JSON (safer than dill for cross-version compatibility)
+            args_json = json.dumps(args, default=str)
+            kwargs_json = json.dumps(kwargs, default=str)
+            
+        except OSError as e:
+            raise Exception(f"Cannot retrieve source code for function {python_func.__name__}: {e}")
+        except Exception as e:
+            raise Exception(f"Failed to reconstruct function source: {e}")
+        
+        # Create Python code that will define and execute the function
+        python_code = f'''
+import json
+import traceback
+
+try:
+    # Define the function from source
+{textwrap.indent(func_source, "    ")}
+    
+    # Deserialize args and kwargs from JSON
+    args_json = """{args_json}"""
+    kwargs_json = """{kwargs_json}"""
+    args = json.loads(args_json)
+    kwargs = json.loads(kwargs_json)
+    
+    # Execute the function
+    result = {func_name}(*args, **kwargs)
+
+    # Create success output payload
+    output_payload = {{
+        "success": True,
+        "result": result,
+        "error": None
+    }}
+    
+except Exception as e:
+    # Create error output payload
+    output_payload = {{
+        "success": False,
+        "result": None,
+        "error": {{
+            "type": type(e).__name__,
+            "message": str(e),
+            "traceback": traceback.format_exc()
+        }}
+    }}
+
+# Serialize the output payload as JSON
+import json
+output_json = json.dumps(output_payload, default=str)
+
+# Print the JSON output with markers
+print(f"<<<VENV_EXEC_START>>>{{output_json}}<<<VENV_EXEC_END>>>")
+'''
+        
+        # Encode the Python code in base64 to avoid shell escaping issues
+        encoded_code = base64.b64encode(python_code.encode('utf-8')).decode('ascii')
+        
+        # Execute the Python code in the virtual environment
+        python_command = f"python -c \"import base64; exec(base64.b64decode('{encoded_code}').decode('utf-8'))\""
+        stdout, stderr = await self.venv_cmd(venv_name, python_command)
+        
+        # Parse the output to extract the payload
+        start_marker = "<<<VENV_EXEC_START>>>"
+        end_marker = "<<<VENV_EXEC_END>>>"
+
+        # Print original stdout
+        print(stdout[:stdout.find(start_marker)])
+        
+        if start_marker in stdout and end_marker in stdout:
+            start_idx = stdout.find(start_marker) + len(start_marker)
+            end_idx = stdout.find(end_marker)
+            
+            if start_idx < end_idx:
+                output_json = stdout[start_idx:end_idx]
+
+                try:
+                    # Decode and deserialize the output payload from JSON
+                    output_payload = json.loads(output_json)
+                except Exception as e:
+                    raise Exception(f"Failed to decode output payload: {e}")
+                
+                if output_payload["success"]:
+                    return output_payload["result"]
+                else:
+                    # Recreate and raise the original exception
+                    error_info = output_payload["error"]
+                    error_class = eval(error_info["type"])
+                    raise error_class(error_info["message"])
+            else:
+                raise Exception("Invalid output format: markers found but no content between them")
+        else:
+            # Fallback: return stdout/stderr if no payload markers found
+            raise Exception(f"No output payload found. stdout: {stdout}, stderr: {stderr}")

tests/pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+asyncio_mode = auto
+markers =
+    asyncio: asyncio mark

tests/venv.py

@@ -0,0 +1,151 @@
+"""
+Virtual Environment Testing Module
+This module tests the ability to execute python code in a virtual environment within C/ua Containers.
+
+Required environment variables:
+- CUA_API_KEY: API key for C/ua cloud provider
+- CUA_CONTAINER_NAME: Name of the container to use
+"""
+
+import os
+import asyncio
+import pytest
+from pathlib import Path
+import sys
+import traceback
+
+# Load environment variables from .env file
+project_root = Path(__file__).parent.parent
+env_file = project_root / ".env"
+print(f"Loading environment from: {env_file}")
+from dotenv import load_dotenv
+
+load_dotenv(env_file)
+
+# Add paths to sys.path if needed
+pythonpath = os.environ.get("PYTHONPATH", "")
+for path in pythonpath.split(":"):
+    if path and path not in sys.path:
+        sys.path.insert(0, path)  # Insert at beginning to prioritize
+        print(f"Added to sys.path: {path}")
+
+from computer.computer import Computer
+from computer.providers.base import VMProviderType
+
+
+@pytest.fixture(scope="session")
+async def computer():
+    """Shared Computer instance for all test cases."""
+    # Create a remote Linux computer with C/ua
+    computer = Computer(
+        os_type="linux",
+        api_key=os.getenv("CUA_API_KEY"),
+        name=str(os.getenv("CUA_CONTAINER_NAME")),
+        provider_type=VMProviderType.CLOUD,
+    )
+    
+    try:
+        await computer.run()
+        yield computer
+    finally:
+        await computer.stop()
+
+
+# Sample test cases
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_install(computer):
+    """Test virtual environment creation and package installation."""
+    # Create a test virtual environment and install requests
+    stdout, _ = await computer.venv_install("test_env", ["requests"])
+    
+    # Check that installation was successful (no major errors)
+    assert "Successfully installed" in stdout or "Requirement already satisfied" in stdout
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_cmd(computer):
+    """Test executing shell commands in virtual environment."""
+    # Test Python version check
+    stdout, _ = await computer.venv_cmd("test_env", "python --version")
+    
+    assert "Python" in stdout
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_exec(computer):
+    """Test executing Python functions in virtual environment."""
+    def test_function(message="Hello World"):
+        import sys
+        return f"Python {sys.version_info.major}.{sys.version_info.minor}: {message}"
+    
+    result = await computer.venv_exec("test_env", test_function, message="Test successful!")
+    
+    assert "Python" in result
+    assert "Test successful!" in result
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_exec_with_package(computer):
+    """Test executing Python functions that use installed packages."""
+    def test_requests():
+        import requests
+        return f"requests version: {requests.__version__}"
+    
+    result = await computer.venv_exec("test_env", test_requests)
+    
+    assert "requests version:" in result
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_exec_error_handling(computer):
+    """Test error handling in venv_exec."""
+    def test_error():
+        raise ValueError("This is a test error")
+    
+    with pytest.raises(ValueError, match="This is a test error"):
+        await computer.venv_exec("test_env", test_error)
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_exec_with_args_kwargs(computer):
+    """Test executing Python functions with args and kwargs that return an object."""
+    def create_data_object(name, age, *hobbies, **metadata):
+        return {
+            "name": name,
+            "age": age,
+            "hobbies": list(hobbies),
+            "metadata": metadata,
+            "status": "active"
+        }
+    
+    args = ["Alice", 25, "reading", "coding"]
+    kwargs = {"location": "New York", "department": "Engineering"}
+
+    result = await computer.venv_exec(
+        "test_env", 
+        create_data_object, 
+        *args, 
+        **kwargs
+    )
+    
+    assert result["name"] == "Alice"
+    assert result["age"] == 25
+    assert result["hobbies"] == ["reading", "coding"]
+    assert result["metadata"]["location"] == "New York"
+    assert result["status"] == "active"
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_venv_exec_stdout_capture(computer, capfd):
+    """Test capturing stdout from Python functions executed in virtual environment."""
+    def hello_world_function():
+        print("Hello World!")
+        return "Function completed"
+    
+    # Execute the function in the virtual environment
+    result = await computer.venv_exec("test_env", hello_world_function)
+    
+    # Capture stdout and stderr
+    out, _ = capfd.readouterr()
+    
+    # Assert the stdout contains our expected output
+    assert out == "Hello World!\n\n"
+    assert result == "Function completed"
+
+if __name__ == "__main__":
+    # Run tests directly
+    pytest.main([__file__, "-v"])