dbx-patch 0.1.0__py3-none-any.whl

dbx_patch/__about__.py

@@ -0,0 +1,2 @@
+# file generated by dynamic versioning during build when enabled
+__version__ = "0.1.0"

dbx_patch/__init__.py

@@ -0,0 +1,41 @@
+from dbx_patch.__about__ import __version__
+from dbx_patch.install_sitecustomize import (
+    check_sitecustomize_status,
+    install_sitecustomize,
+    uninstall_sitecustomize,
+)
+from dbx_patch.models import (
+    ApplyPatchesResult,
+    PatchResult,
+    PthProcessingResult,
+    RemovePatchesResult,
+    SitecustomizeStatus,
+    StatusResult,
+    VerifyResult,
+)
+from dbx_patch.patch_dbx import (
+    check_patch_status,
+    patch_and_install,
+    patch_dbx,
+    remove_all_patches,
+    verify_editable_installs,
+)
+
+__all__ = [
+    "__version__",
+    "check_patch_status",
+    "remove_all_patches",
+    "verify_editable_installs",
+    "install_sitecustomize",
+    "uninstall_sitecustomize",
+    "check_sitecustomize_status",
+    "patch_dbx",
+    "patch_and_install",
+    "ApplyPatchesResult",
+    "PatchResult",
+    "PthProcessingResult",
+    "RemovePatchesResult",
+    "SitecustomizeStatus",
+    "StatusResult",
+    "VerifyResult",
+]

dbx_patch/base_patch.py

@@ -0,0 +1,223 @@
+"""Base classes for DBX-Patch interface.
+
+Provides abstract base classes and singleton metaclass for all patches.
+"""
+
+from abc import ABCMeta, abstractmethod
+import contextlib
+import threading
+from typing import Any
+
+from dbx_patch.models import PatchResult
+
+
+class SingletonMeta(ABCMeta):
+    """Thread-safe singleton metaclass for patch classes.
+
+    Ensures only one instance of each patch class exists, with thread-safe
+    creation in notebook environments. Inherits from ABCMeta to support ABC.
+    """
+
+    _instances: dict[type, Any] = {}
+    _lock: threading.Lock = threading.Lock()
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        """Get or create the singleton instance."""
+        if cls not in cls._instances:
+            with cls._lock:
+                # Double-check pattern
+                if cls not in cls._instances:
+                    instance = super().__call__(*args, **kwargs)
+                    cls._instances[cls] = instance
+        return cls._instances[cls]
+
+    @classmethod
+    def reset_instance(mcs, cls: type) -> None:
+        """Reset singleton instance for testing.
+
+        Args:
+            cls: The class whose singleton instance should be reset
+        """
+        with mcs._lock:
+            if cls in mcs._instances:
+                del mcs._instances[cls]
+
+
+class BasePatch(metaclass=SingletonMeta):
+    """Abstract base class for all Databricks runtime patches.
+
+    Provides a unified interface for applying, removing, and checking patch status.
+    All patches follow the singleton pattern to ensure consistent global state.
+
+    Attributes:
+        _is_applied: Whether the patch has been applied
+        _original_target: Reference to original function/method for restoration
+        _cached_editable_paths: Set of cached editable install paths (if applicable)
+        _logger: Cached logger instance
+    """
+
+    def __init__(self, verbose: bool = True) -> None:
+        """Initialize the patch (called only once due to singleton).
+
+        Args:
+            verbose: Enable verbose logging
+        """
+        self._is_applied: bool = False
+        self._original_target: Any = None
+        self._cached_editable_paths: set[str] = set()
+        self._verbose: bool = verbose
+        self._logger: Any = None
+
+    def _get_logger(self) -> Any:
+        """Get lazily-initialized logger instance.
+
+        Returns:
+            Logger instance or None if unavailable
+        """
+        if self._logger is None:
+            with contextlib.suppress(Exception):
+                from dbx_patch.utils.logger import get_logger
+
+                self._logger = get_logger()
+        return self._logger
+
+    def _detect_editable_paths(self) -> set[str]:
+        """Detect editable install paths from pth_processor.
+
+        Returns:
+            Set of absolute paths to editable install directories
+        """
+        try:
+            from dbx_patch.pth_processor import get_editable_install_paths
+
+            return get_editable_install_paths()
+        except Exception:
+            return set()
+
+    @abstractmethod
+    def patch(self) -> PatchResult:
+        """Apply the patch to the Databricks runtime.
+
+        This method should:
+        1. Check if already applied (return early if so)
+        2. Verify target modules/functions exist
+        3. Apply the patch (monkey-patch, wrap, register, etc.)
+        4. Update internal state (_is_applied, _original_target, etc.)
+        5. Return PatchResult with operation details
+
+        Returns:
+            PatchResult indicating success/failure and details
+        """
+        ...
+
+    @abstractmethod
+    def remove(self) -> bool:
+        """Remove the patch and restore original behavior.
+
+        This method should:
+        1. Check if patch is applied (return early if not)
+        2. Restore original functions/methods from _original_target
+        3. Clean up any registered callbacks or hooks
+        4. Reset internal state (_is_applied = False, etc.)
+        5. Return True if successful, False otherwise
+
+        Returns:
+            True if patch was removed successfully, False otherwise
+        """
+        ...
+
+    @abstractmethod
+    def is_applied(self) -> bool:
+        """Check if the patch is currently applied.
+
+        Returns:
+            True if patch is applied, False otherwise
+        """
+        ...
+
+    def refresh_paths(self) -> int:
+        """Refresh cached editable install paths (optional, override if needed).
+
+        Returns:
+            Number of editable paths detected
+        """
+        self._cached_editable_paths = self._detect_editable_paths()
+        return len(self._cached_editable_paths)
+
+    def get_editable_paths(self) -> set[str]:
+        """Get current cached editable paths (optional, override if needed).
+
+        Returns:
+            Set of editable install paths
+        """
+        return self._cached_editable_paths.copy()
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton instance for testing purposes."""
+        SingletonMeta.reset_instance(cls)
+
+
+class BaseVerification(metaclass=SingletonMeta):
+    """Abstract base class for verification-only patches.
+
+    Verification patches don't modify runtime behavior - they only check
+    compatibility and report findings.
+
+    Attributes:
+        _is_verified: Whether verification has been performed
+        _logger: Cached logger instance
+    """
+
+    def __init__(self, verbose: bool = True) -> None:
+        """Initialize the verification (called only once due to singleton).
+
+        Args:
+            verbose: Enable verbose logging
+        """
+        self._is_verified: bool = False
+        self._verbose: bool = verbose
+        self._logger: Any = None
+
+    def _get_logger(self) -> Any:
+        """Get lazily-initialized logger instance.
+
+        Returns:
+            Logger instance or None if unavailable
+        """
+        if self._logger is None:
+            with contextlib.suppress(Exception):
+                from dbx_patch.utils.logger import get_logger
+
+                self._logger = get_logger()
+        return self._logger
+
+    @abstractmethod
+    def verify(self) -> PatchResult:
+        """Verify compatibility with Databricks runtime.
+
+        This method should:
+        1. Check if already verified (return early if so)
+        2. Import and inspect target modules/hooks
+        3. Verify they won't interfere with editable installs
+        4. Update internal state (_is_verified = True)
+        5. Return PatchResult with verification details
+
+        Returns:
+            PatchResult indicating compatibility status
+        """
+        ...
+
+    @abstractmethod
+    def is_verified(self) -> bool:
+        """Check if verification has been performed.
+
+        Returns:
+            True if verified, False otherwise
+        """
+        ...
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton instance for testing purposes."""
+        SingletonMeta.reset_instance(cls)

dbx_patch/cli.py

@@ -0,0 +1,44 @@
+import contextlib
+
+from dbx_patch.__about__ import __version__
+from dbx_patch.patch_dbx import check_patch_status, patch_dbx, remove_all_patches, verify_editable_installs
+
+# Module-level logger
+_logger = None
+with contextlib.suppress(Exception):
+    from dbx_patch.utils.logger import get_logger
+
+    _logger = get_logger()
+
+
+def main() -> None:
+    if _logger:
+        _logger.info(f"dbx-patch v{__version__}!")
+
+    # Allow running as a script
+    import argparse
+
+    parser = argparse.ArgumentParser(description="DBX-Patch for editable installs")
+    parser.add_argument("--apply", action="store_true", help="Apply all patches")
+    parser.add_argument("--verify", action="store_true", help="Verify configuration")
+    parser.add_argument("--status", action="store_true", help="Check patch status")
+    parser.add_argument("--remove", action="store_true", help="Remove all patches")
+    parser.add_argument("--quiet", action="store_true", help="Suppress output")
+
+    args = parser.parse_args()
+
+    if args.apply:
+        patch_dbx()
+    elif args.verify:
+        verify_editable_installs()
+    elif args.status:
+        check_patch_status()
+    elif args.remove:
+        remove_all_patches()
+    else:
+        # Default: apply patches
+        patch_dbx()
+
+
+if __name__ == "__main__":
+    main()

dbx_patch/install_sitecustomize.py

@@ -0,0 +1,286 @@
+"""Install sitecustomize.py to Auto-Apply Patches on Startup.
+
+This module provides functionality to install a sitecustomize.py file that
+automatically applies all dbx-patch fixes when Python starts. This solves
+the timing issue where sys_path_init and WsfsImportHook are loaded before
+any notebook code runs.
+
+Why sitecustomize.py?
+- Python automatically imports sitecustomize.py during interpreter initialization
+- It runs BEFORE sys_path_init.py and import hooks are installed
+- This is the ONLY way to patch the import system early enough
+
+Usage:
+    from dbx_patch.install_sitecustomize import install_sitecustomize
+    install_sitecustomize()
+"""
+
+from pathlib import Path
+import sys
+
+from dbx_patch.models import SitecustomizeStatus
+from dbx_patch.utils.logger import PatchLogger
+
+logger = PatchLogger()
+
+
+def get_site_packages_path() -> Path | None:
+    """Get the first writable site-packages directory.
+
+    Returns:
+        Path to site-packages directory, or None if not found
+    """
+    for path_str in sys.path:
+        if "site-packages" in path_str:
+            path = Path(path_str)
+            if path.exists() and path.is_dir():
+                # Check if writable
+                try:
+                    test_file = path / ".dbx_patch_write_test"
+                    test_file.touch()
+                    test_file.unlink()
+                    return path
+                except (OSError, PermissionError):
+                    continue
+    return None
+
+
+def get_sitecustomize_content() -> str:
+    """Generate the sitecustomize.py content.
+
+    Returns:
+        Python code to be written to sitecustomize.py
+    """
+    return """\"\"\"Auto-apply dbx-patch on Python startup.
+
+This file is automatically loaded by Python during interpreter initialization.
+It applies all dbx-patch fixes BEFORE sys_path_init and import hooks are loaded.
+
+Generated by: dbx-patch
+DO NOT EDIT MANUALLY - Use dbx_patch.install_sitecustomize() to update
+\"\"\"
+
+import sys
+
+
+def _apply_dbx_patch() -> None:
+    \"\"\"Apply dbx-patch fixes silently during startup.\"\"\"
+    try:
+        # Import and apply all patches
+        from dbx_patch.patch_dbx import patch_dbx
+
+        # Apply silently (no output to avoid cluttering startup)
+        patch_dbx(force_refresh=False)
+
+    except ImportError:
+        # dbx-patch not installed, skip silently
+        pass
+    except Exception as e:
+        # Log error but don't break Python startup
+        print(f"Warning: dbx-patch auto-apply failed: {e}", file=sys.stderr)
+
+
+# Apply patches immediately on import
+_apply_dbx_patch()
+"""
+
+
+def install_sitecustomize(force: bool = True, restart_python: bool = True) -> bool:
+    """Install sitecustomize.py to auto-apply patches on Python startup.
+
+    This is the RECOMMENDED way to use dbx-patch because:
+    1. Patches are applied BEFORE sys_path_init runs
+    2. Import hooks are patched BEFORE they're installed
+    3. No need to manually call patch_dbx() in every notebook
+    4. Works automatically for all Python processes on the cluster
+
+    Args:
+        force: If True, overwrite existing sitecustomize.py
+        restart_python: If True, automatically restart Python using dbutils.library.restartPython()
+
+    Returns:
+        True if installation succeeded, False otherwise
+
+    Example:
+        # Run once per cluster (e.g., in init script or first notebook):
+        from dbx_patch.install_sitecustomize import install_sitecustomize
+        install_sitecustomize()
+
+        # Python will restart automatically if running in Databricks
+        # After restart, editable installs will work automatically!
+    """
+    with logger.section("Installing sitecustomize.py for auto-apply"):
+        # Find site-packages
+        site_packages = get_site_packages_path()
+        if site_packages is None:
+            logger.error("Could not find writable site-packages directory")
+            with logger.indent():
+                logger.info("Make sure you have write permissions to site-packages")
+            return False
+
+        sitecustomize_path = site_packages / "sitecustomize.py"
+
+        # Check if already exists
+        if sitecustomize_path.exists() and not force:
+            logger.warning(f"sitecustomize.py already exists: {sitecustomize_path}")
+            with logger.indent():
+                logger.info("Use force=True to overwrite")
+                logger.info("Or manually merge the content")
+            return False
+
+        # Backup existing file if it exists
+        if sitecustomize_path.exists():
+            backup_path = site_packages / "sitecustomize.py.backup"
+            logger.info(f"Backing up existing file to: {backup_path}")
+            try:
+                sitecustomize_path.rename(backup_path)
+            except OSError as e:
+                logger.error(f"Failed to backup existing file: {e}")  # noqa: TRY400
+                return False
+
+        # Write new sitecustomize.py
+        try:
+            content = get_sitecustomize_content()
+            sitecustomize_path.write_text(content, encoding="utf-8")
+
+            logger.success(f"sitecustomize.py installed: {sitecustomize_path}")
+            logger.blank()
+            logger.info("✅ Installation complete!")
+
+            # Try to restart Python automatically if in Databricks environment
+            if restart_python:
+                logger.blank()
+                logger.info("Attempting to restart Python kernel...")
+
+                try:
+                    # Try to access dbutils (available in Databricks notebooks)
+                    # dbutils is injected by Databricks and available as a variable
+                    try:
+                        logger.info("Restarting Python kernel via dbutils.library.restartPython()...")
+
+                        dbutils.library.restartPython()  # ty:ignore[unresolved-reference]  # noqa: F821
+
+                    except Exception:
+                        # Not in Databricks environment
+                        logger.blank()
+                        logger.warning("Not running in Databricks environment")
+                        logger.info("Next steps:")
+                        with logger.indent():
+                            logger.info("1. Restart your Python kernel/notebook manually")
+                            logger.info("2. Editable installs will work automatically")
+                            logger.info("3. No need to call patch_dbx() anymore!")
+
+                except Exception as e:
+                    # Failed to restart, provide manual instructions
+                    logger.blank()
+                    logger.warning(f"Could not restart Python automatically: {e}")
+                    logger.info("Next steps:")
+                    with logger.indent():
+                        logger.info("1. Restart your Python kernel/notebook manually")
+                        logger.info("2. Editable installs will work automatically")
+                        logger.info("3. No need to call patch_dbx() anymore!")
+            else:
+                logger.blank()
+                logger.info("Next steps:")
+                with logger.indent():
+                    logger.info("1. Restart your Python kernel/notebook")
+                    logger.info("2. Editable installs will work automatically")
+                    logger.info("3. No need to call patch_dbx() anymore!")
+
+            return True
+
+        except OSError as e:
+            logger.error(f"Failed to write sitecustomize.py: {e}")  # noqa: TRY400
+            return False
+
+
+def uninstall_sitecustomize() -> bool:
+    """Remove the auto-apply sitecustomize.py.
+
+    Returns:
+        True if uninstallation succeeded, False otherwise
+    """
+    logger.debug("uninstall_sitecustomize() called")
+
+    with logger.section("Uninstalling sitecustomize.py"):
+        site_packages = get_site_packages_path()
+        if site_packages is None:
+            logger.error("Could not find site-packages directory")
+            return False
+
+        sitecustomize_path = site_packages / "sitecustomize.py"
+
+        if not sitecustomize_path.exists():
+            logger.info("sitecustomize.py does not exist, nothing to uninstall")
+            return True
+
+        # Check if it's our file
+        try:
+            content = sitecustomize_path.read_text(encoding="utf-8")
+            if "dbx-patch" not in content and "dbx_patch" not in content:
+                logger.warning("sitecustomize.py exists but wasn't created by dbx-patch")
+                with logger.indent():
+                    logger.info("Skipping removal for safety")
+                    logger.info("Manual removal required if needed")
+                return False
+        except OSError as e:
+            logger.error(f"Failed to read sitecustomize.py: {e}")  # noqa: TRY400
+            return False
+
+        # Remove the file
+        try:
+            sitecustomize_path.unlink()
+            logger.success("sitecustomize.py removed")
+
+            # Restore backup if it exists
+            backup_path = site_packages / "sitecustomize.py.backup"
+            if backup_path.exists():
+                backup_path.rename(sitecustomize_path)
+                logger.info("Restored backup file")
+
+            return True
+
+        except OSError as e:
+            logger.error(f"Failed to remove sitecustomize.py: {e}")  # noqa: TRY400
+            return False
+
+
+def check_sitecustomize_status() -> SitecustomizeStatus:
+    """Check if sitecustomize.py is installed and active.
+
+    Returns:
+        SitecustomizeStatus with installation information
+    """
+    logger.debug("check_sitecustomize_status() called")
+
+    site_packages = get_site_packages_path()
+    if site_packages is None:
+        logger.warning("Could not find site-packages directory")
+        return SitecustomizeStatus(
+            installed=False,
+            path=None,
+            is_dbx_patch=False,
+        )
+
+    sitecustomize_path = site_packages / "sitecustomize.py"
+    installed = sitecustomize_path.exists()
+    is_dbx_patch = False
+
+    if installed:
+        try:
+            content = sitecustomize_path.read_text(encoding="utf-8")
+            is_dbx_patch = "dbx-patch" in content or "dbx_patch" in content
+        except OSError:
+            pass
+
+    logger.info("sitecustomize.py status:")
+    with logger.indent():
+        logger.info(f"Installed: {installed}")
+        logger.info(f"Path: {sitecustomize_path}")
+        logger.info(f"Created by dbx-patch: {is_dbx_patch}")
+
+    return SitecustomizeStatus(
+        installed=installed,
+        path=str(sitecustomize_path) if sitecustomize_path else None,
+        is_dbx_patch=is_dbx_patch,
+    )

dbx_patch/models.py

@@ -0,0 +1,94 @@
+"""Data models for DBX-Patch results and status.
+
+Provides strongly-typed dataclasses for all function return values.
+"""
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class PthProcessingResult:
+    """Results from processing .pth files."""
+
+    site_dirs_scanned: int
+    pth_files_found: int
+    paths_extracted: list[str]
+    egg_link_paths: list[str]
+    metadata_paths: list[str]
+    paths_added: int
+    total_editable_paths: int
+
+
+@dataclass
+class PatchResult:
+    """Generic patch operation result."""
+
+    success: bool
+    already_patched: bool
+    function_found: bool = True
+    hook_found: bool = True
+    editable_paths_count: int = 0
+    editable_paths: list[str] = field(default_factory=list)
+    error: str | None = None
+
+
+@dataclass
+class ApplyPatchesResult:
+    """Results from applying all patches."""
+
+    sys_path_init_patch: PatchResult | None
+    pth_processing: PthProcessingResult | None
+    wsfs_hook_patch: PatchResult | None
+    wsfs_path_finder_patch: PatchResult | None
+    python_path_hook_patch: PatchResult | None
+    autoreload_hook_patch: PatchResult | None
+    overall_success: bool
+    editable_paths: list[str]
+
+
+@dataclass
+class VerifyResult:
+    """Results from verifying editable install configuration."""
+
+    editable_paths: list[str]
+    paths_in_sys_path: list[str]
+    wsfs_hook_patched: bool
+    wsfs_path_finder_patched: bool
+    python_path_hook_patched: bool
+    autoreload_hook_patched: bool
+    importable_packages: list[str]
+    status: str  # 'ok', 'warning', or 'error'
+
+
+@dataclass
+class StatusResult:
+    """Current patch status."""
+
+    sys_path_init_patched: bool
+    wsfs_hook_patched: bool
+    wsfs_path_finder_patched: bool
+    python_path_hook_patched: bool
+    autoreload_hook_patched: bool
+    editable_paths_count: int
+    pth_files_processed: bool
+
+
+@dataclass
+class RemovePatchesResult:
+    """Results from removing all patches."""
+
+    sys_path_init_unpatched: bool
+    wsfs_hook_unpatched: bool
+    wsfs_path_finder_unpatched: bool
+    python_path_hook_unpatched: bool
+    autoreload_hook_unpatched: bool
+    success: bool
+
+
+@dataclass
+class SitecustomizeStatus:
+    """Status of sitecustomize.py installation."""
+
+    installed: bool
+    path: str | None
+    is_dbx_patch: bool