mlenvdoctor 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

mlenvdoctor/utils.py

@@ -5,11 +5,20 @@ import sys
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import sys
+
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
-from rich.text import Text
 
-console = Console()
+from .exceptions import DiagnosticError
+from .icons import icon_check, icon_cross, icon_info, icon_warning
+
+# Configure console for Windows compatibility
+if sys.platform == "win32":
+    # Use legacy Windows renderer if needed
+    console = Console(legacy_windows=True, force_terminal=True)
+else:
+    console = Console()
 
 
 def run_command(
@@ -17,8 +26,29 @@ def run_command(
     capture_output: bool = True,
     check: bool = False,
     timeout: Optional[int] = 30,
-) -> subprocess.CompletedProcess:
-    """Run a shell command with error handling."""
+) -> subprocess.CompletedProcess[str]:
+    """
+    Run a shell command with error handling and input validation.
+
+    Args:
+        cmd: Command and arguments as a list
+        capture_output: Whether to capture stdout/stderr
+        check: Whether to raise on non-zero exit code
+        timeout: Command timeout in seconds
+
+    Returns:
+        CompletedProcess with command result
+
+    Raises:
+        DiagnosticError: For command execution errors
+        ConfigurationError: For invalid input
+    """
+    from .validators import sanitize_command, validate_timeout
+
+    # Validate and sanitize inputs
+    cmd = sanitize_command(cmd)
+    timeout = validate_timeout(timeout)
+
     try:
         result = subprocess.run(
             cmd,
@@ -28,29 +58,58 @@ def run_command(
             timeout=timeout,
         )
         return result
-    except subprocess.TimeoutExpired:
-        console.print(f"[red]Command timed out: {' '.join(cmd)}[/red]")
-        raise
-    except FileNotFoundError:
-        console.print(f"[red]Command not found: {cmd[0]}[/red]")
-        raise
+    except subprocess.TimeoutExpired as e:
+        error_msg = f"Command timed out after {timeout}s: {' '.join(cmd)}"
+        console.print(f"[red]{error_msg}[/red]")
+        raise DiagnosticError(
+            error_msg,
+            "Try increasing timeout or check if the command is hanging",
+        ) from e
+    except FileNotFoundError as e:
+        error_msg = f"Command not found: {cmd[0]}"
+        console.print(f"[red]{error_msg}[/red]")
+        raise DiagnosticError(
+            error_msg,
+            f"Install {cmd[0]} or ensure it's in your PATH",
+        ) from e
     except subprocess.CalledProcessError as e:
         if not check:
-            return e  # type: ignore
+            # Return the exception as if it were a result
+            # This maintains backward compatibility but is type-unsafe
+            return subprocess.CompletedProcess(  # type: ignore[return-value]
+                cmd, e.returncode, e.stdout, e.stderr
+            )
         raise
 
 
 def check_command_exists(cmd: str) -> bool:
-    """Check if a command exists in PATH."""
+    """
+    Check if a command exists in PATH.
+
+    Args:
+        cmd: Command name to check
+
+    Returns:
+        True if command exists and is executable, False otherwise
+    """
+    if not isinstance(cmd, str) or not cmd.strip():
+        return False
+
     try:
-        subprocess.run(
-            [cmd, "--version"] if cmd != "nvidia-smi" else [cmd],
+        # Use 'which' on Unix, 'where' on Windows
+        if sys.platform == "win32":
+            check_cmd = ["where", cmd]
+        else:
+            check_cmd = ["which", cmd]
+
+        result = subprocess.run(
+            check_cmd,
             capture_output=True,
             timeout=5,
             check=False,
         )
-        return True
-    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
         return False
 
 
@@ -64,22 +123,22 @@ def get_home_config_dir() -> Path:
 
 def print_success(message: str) -> None:
     """Print a success message."""
-    console.print(f"[green]✅ {message}[/green]")
+    console.print(f"[green]{icon_check()} {message}[/green]")
 
 
 def print_error(message: str) -> None:
     """Print an error message."""
-    console.print(f"[red]❌ {message}[/red]")
+    console.print(f"[red]{icon_cross()} {message}[/red]")
 
 
 def print_warning(message: str) -> None:
     """Print a warning message."""
-    console.print(f"[yellow]⚠️  {message}[/yellow]")
+    console.print(f"[yellow]{icon_warning()}  {message}[/yellow]")
 
 
 def print_info(message: str) -> None:
     """Print an info message."""
-    console.print(f"[blue]ℹ️  {message}[/blue]")
+    console.print(f"[blue]{icon_info()}  {message}[/blue]")
 
 
 def with_spinner(message: str):
@@ -103,5 +162,3 @@ def format_size(size_bytes: int) -> str:
 def get_python_version() -> Tuple[int, int, int]:
     """Get Python version as tuple."""
     return sys.version_info[:3]
-
-

mlenvdoctor/validators.py

@@ -0,0 +1,217 @@
+"""Input validation and sanitization for ML Environment Doctor."""
+
+import re
+from pathlib import Path
+from typing import Optional
+
+from .exceptions import ConfigurationError
+
+
+def validate_model_name(model_name: str) -> str:
+    """
+    Validate and sanitize model name.
+
+    Args:
+        model_name: Model name to validate
+
+    Returns:
+        Sanitized model name
+
+    Raises:
+        ConfigurationError: If model name is invalid
+    """
+    if not model_name or not isinstance(model_name, str):
+        raise ConfigurationError(
+            "Model name must be a non-empty string",
+            "Use a valid model name like 'tinyllama', 'gpt2', or 'mistral-7b'",
+        )
+
+    # Remove whitespace
+    model_name = model_name.strip()
+
+    # Check for dangerous characters (basic sanitization)
+    if not re.match(r"^[a-zA-Z0-9._-]+$", model_name):
+        raise ConfigurationError(
+            f"Invalid model name: {model_name}",
+            "Model name can only contain letters, numbers, dots, underscores, and hyphens",
+        )
+
+    return model_name.lower()
+
+
+def validate_file_path(file_path: Path, must_exist: bool = False, must_be_file: bool = False) -> Path:
+    """
+    Validate and sanitize file path.
+
+    Args:
+        file_path: Path to validate
+        must_exist: Whether the path must exist
+        must_be_file: Whether the path must be a file
+
+    Returns:
+        Resolved, absolute path
+
+    Raises:
+        ConfigurationError: If path is invalid
+    """
+    if not isinstance(file_path, (Path, str)):
+        raise ConfigurationError(
+            "File path must be a Path object or string",
+            "Use pathlib.Path or a valid string path",
+        )
+
+    path = Path(file_path).resolve()
+
+    # Check for path traversal attempts
+    if ".." in str(path):
+        # Resolve should handle this, but double-check
+        resolved = path.resolve()
+        if ".." in str(resolved):
+            raise ConfigurationError(
+                "Invalid path: contains '..'",
+                "Use absolute paths or relative paths without '..'",
+            )
+
+    if must_exist and not path.exists():
+        raise ConfigurationError(
+            f"Path does not exist: {path}",
+            "Ensure the file or directory exists",
+        )
+
+    if must_be_file and not path.is_file():
+        raise ConfigurationError(
+            f"Path is not a file: {path}",
+            "Provide a valid file path",
+        )
+
+    return path
+
+
+def validate_log_level(level: str) -> str:
+    """
+    Validate logging level.
+
+    Args:
+        level: Log level to validate
+
+    Returns:
+        Validated log level
+
+    Raises:
+        ConfigurationError: If level is invalid
+    """
+    valid_levels = {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"}
+    level_upper = level.upper() if isinstance(level, str) else str(level).upper()
+
+    if level_upper not in valid_levels:
+        raise ConfigurationError(
+            f"Invalid log level: {level}",
+            f"Use one of: {', '.join(valid_levels)}",
+        )
+
+    return level_upper
+
+
+def validate_stack_name(stack: str) -> str:
+    """
+    Validate ML stack name.
+
+    Args:
+        stack: Stack name to validate
+
+    Returns:
+        Validated stack name
+
+    Raises:
+        ConfigurationError: If stack is invalid
+    """
+    valid_stacks = {"trl-peft", "minimal"}
+    stack_lower = stack.lower() if isinstance(stack, str) else str(stack).lower()
+
+    if stack_lower not in valid_stacks:
+        raise ConfigurationError(
+            f"Invalid stack: {stack}",
+            f"Use one of: {', '.join(valid_stacks)}",
+        )
+
+    return stack_lower
+
+
+def sanitize_command(cmd: list[str]) -> list[str]:
+    """
+    Sanitize command arguments to prevent injection.
+
+    Args:
+        cmd: Command and arguments list
+
+    Returns:
+        Sanitized command list
+
+    Raises:
+        ConfigurationError: If command contains dangerous patterns
+    """
+    if not isinstance(cmd, list) or not cmd:
+        raise ConfigurationError(
+            "Command must be a non-empty list",
+            "Provide command as a list of strings",
+        )
+
+    sanitized = []
+    for arg in cmd:
+        if not isinstance(arg, str):
+            raise ConfigurationError(
+                "All command arguments must be strings",
+                "Convert all arguments to strings",
+            )
+
+        # Check for command injection patterns
+        dangerous_patterns = [";", "&&", "||", "`", "$(", "<", ">", "|"]
+        for pattern in dangerous_patterns:
+            if pattern in arg:
+                raise ConfigurationError(
+                    f"Dangerous pattern detected in command: {pattern}",
+                    "Do not use shell operators in command arguments",
+                )
+
+        sanitized.append(arg)
+
+    return sanitized
+
+
+def validate_timeout(timeout: Optional[int], min_timeout: int = 1, max_timeout: int = 3600) -> Optional[int]:
+    """
+    Validate timeout value.
+
+    Args:
+        timeout: Timeout in seconds
+        min_timeout: Minimum allowed timeout
+        max_timeout: Maximum allowed timeout
+
+    Returns:
+        Validated timeout
+
+    Raises:
+        ConfigurationError: If timeout is invalid
+    """
+    if timeout is None:
+        return None
+
+    if not isinstance(timeout, int):
+        raise ConfigurationError(
+            f"Timeout must be an integer, got {type(timeout)}",
+            "Provide timeout as an integer number of seconds",
+        )
+
+    if timeout < min_timeout:
+        raise ConfigurationError(
+            f"Timeout too small: {timeout}s (minimum: {min_timeout}s)",
+            f"Increase timeout to at least {min_timeout} seconds",
+        )
+
+    if timeout > max_timeout:
+        raise ConfigurationError(
+            f"Timeout too large: {timeout}s (maximum: {max_timeout}s)",
+            f"Decrease timeout to at most {max_timeout} seconds",
+        )
+
+    return timeout

{mlenvdoctor-0.1.0.dist-info → mlenvdoctor-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlenvdoctor
-Version: 0.1.0
+Version: 0.1.2
 Summary: Diagnose & fix ML environments for LLM fine-tuning
 Author: ML Environment Doctor Contributors
 License: MIT
@@ -20,6 +20,7 @@ Requires-Python: >=3.8
 Requires-Dist: packaging>=23.0
 Requires-Dist: psutil>=5.9.0
 Requires-Dist: rich>=13.0.0
+Requires-Dist: tomli>=2.0.0; python_version < '3.11'
 Requires-Dist: typer>=0.9.0
 Provides-Extra: dev
 Requires-Dist: black>=23.0.0; extra == 'dev'
@@ -34,7 +35,7 @@ Description-Content-Type: text/markdown
 
 [![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![PyPI version](https://badge.fury.io/py/mlenvdoctor.svg)](https://badge.fury.io/py/mlenvdoctor)
+[![PyPI](https://img.shields.io/pypi/v/mlenvdoctor.svg)]([https://pypi.org/project/mlenvdoctor/])
 
 > **Single command fixes 90% of "my torch.cuda.is_available() is False" issues.**
 

mlenvdoctor-0.1.2.dist-info/RECORD

@@ -0,0 +1,21 @@
+mlenvdoctor/__init__.py,sha256=igtJXQ-DiuG4_2BcfW64KPxvhCzSeuXUNUVqDzuLDI8,327
+mlenvdoctor/cli.py,sha256=jvc1MWNciyRGedYCEKSCZeqXfX4wf6xe38Vx44k86xc,7145
+mlenvdoctor/config.py,sha256=c7WUd8XUAo8nz-Ri4TIldc-_p3wiyDNqOsL-MhV9jbI,4770
+mlenvdoctor/constants.py,sha256=ZoYf6-gqmw88CGENyYVur_hCmdcUm0IR5N9pUdOLjJY,1980
+mlenvdoctor/diagnose.py,sha256=7lsAVOPeK2I3nmlzMKmtniyEh1UGv76bJqlqSpjvPL8,20553
+mlenvdoctor/dockerize.py,sha256=AC8HX5sRkSFAM0O0caBnKW4HAdS49MVmMcsplKEDXI4,5562
+mlenvdoctor/exceptions.py,sha256=8wzZE-In0zimXJ3omUA3YmFeklcseO53kqo0SpB58DM,1069
+mlenvdoctor/export.py,sha256=CsuLpbpR2OVs-bud27K8Xv28KCI9veNPIWCsznyFmaw,8671
+mlenvdoctor/fix.py,sha256=fXS4uxBN-FWRFKowmOsdPYI7bnY8jnYxG7UADpJ1hwc,8989
+mlenvdoctor/gpu.py,sha256=sMFgtF4pt-dpOr6IDxvm6f0ChfmCf58E8mApIH6jvAs,6295
+mlenvdoctor/icons.py,sha256=vu35SuBxlZu55rUqldgQfX4UeHeENt082ighZKnzHZY,2289
+mlenvdoctor/logger.py,sha256=OKJQjcdOspARokBcIDyCri14gS_7MlrCt2B3znKc34Y,2377
+mlenvdoctor/parallel.py,sha256=HVJmu8t4k2-XeoQPHlLmhInLOcXjUTiP0VM3THOkAmE,3630
+mlenvdoctor/retry.py,sha256=ZH-KWe6BPNK1sUNFSK2uwdobZ_Z77fxMUfSrSpInt3c,2993
+mlenvdoctor/utils.py,sha256=ehohh-iRLe2qkOMxj5v9yTWONf5gWSdY6CvfrRttTlg,4862
+mlenvdoctor/validators.py,sha256=Kz1FcJM4Cym-S_z5vTocv0cxzKOEgvZIqv8C8c1gSzY,6109
+mlenvdoctor-0.1.2.dist-info/METADATA,sha256=7Rlpv9kjMHQWPBsnZR1gbaUzocWNXNJ3wdGy5UXaNQU,8942
+mlenvdoctor-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mlenvdoctor-0.1.2.dist-info/entry_points.txt,sha256=Y-WH-ANeiTdECIaqi_EB3ZEf_kACkvsYBHnNhXsCI4k,52
+mlenvdoctor-0.1.2.dist-info/licenses/LICENSE,sha256=rGHdyWGvGWYnEFlthqtB-RtRCTa7WaAOElom5qD-nHw,1114
+mlenvdoctor-0.1.2.dist-info/RECORD,,

mlenvdoctor-0.1.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-mlenvdoctor/__init__.py,sha256=hPyn5R1E9SDxyqhMb9AUsDw-TRxj9K2AWcadl5LmvqY,112
-mlenvdoctor/cli.py,sha256=DjF6UYC4QlrGSvbmG3wokX0urAWls9hu0FSIh-jZyYA,5305
-mlenvdoctor/diagnose.py,sha256=zk_E3UF2OlJr6ZFx3OjejyswgcZWX7qz1BHDXCB-vmk,16980
-mlenvdoctor/dockerize.py,sha256=q7afAUSkpL3RHSmbMZy2G9VZ7yRbus4xDADbSjIIuJM,5615
-mlenvdoctor/fix.py,sha256=V-mK30r2xSk-_3uuEHqH01iA9Vt8mjjlCisb4kL_A_Q,8959
-mlenvdoctor/gpu.py,sha256=PTS_dj6JaAoXKHQNzqgh8xOuZHomPoQxuLQMGWQXHqQ,6239
-mlenvdoctor/utils.py,sha256=fAZHgVX6iyOpcd4NU-oSyLjFMO1AXnfgOL5_KL_-0Po,2925
-mlenvdoctor-0.1.0.dist-info/METADATA,sha256=mZnlvzRuaaW9ydTZf5HsP4-wtDs7LUKUEde2xu_STmo,8889
-mlenvdoctor-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mlenvdoctor-0.1.0.dist-info/entry_points.txt,sha256=Y-WH-ANeiTdECIaqi_EB3ZEf_kACkvsYBHnNhXsCI4k,52
-mlenvdoctor-0.1.0.dist-info/licenses/LICENSE,sha256=rGHdyWGvGWYnEFlthqtB-RtRCTa7WaAOElom5qD-nHw,1114
-mlenvdoctor-0.1.0.dist-info/RECORD,,