import_or_install.py

from typing import TYPE_CHECKING, Optional


def import_or_install(package: str, submodule: Optional[str] = None):
    try:
        if submodule:
            module = __import__(f"{package}.{submodule}", fromlist=[submodule])
        else:
            module = __import__(package)
        return module
    except ImportError:
        print(f"Installing {package}...")
        import subprocess
        import sys

        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
        if submodule:
            module = __import__(f"{package}.{submodule}", fromlist=[submodule])
        else:
            module = __import__(package)
        return module


prompt = import_or_install("rich", "prompt")  # type: ignore


if TYPE_CHECKING: # for type checking
    from rich import prompt