83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import subprocess
|
|
from functools import lru_cache
|
|
from logging import getLogger
|
|
from typing import Callable, TypeVar
|
|
|
|
from ...import_utils import patch_object
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
__all__ = ["require_jupyter_kernel_gateway_installed", "skip_on_missing_jupyter_kernel_gateway"]
|
|
|
|
|
|
@lru_cache
|
|
def is_jupyter_kernel_gateway_installed() -> bool:
|
|
"""Check if jupyter-kernel-gateway is installed."""
|
|
try:
|
|
subprocess.run(
|
|
["jupyter", "kernelgateway", "--version"],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=True,
|
|
)
|
|
return True
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
logger.warning(
|
|
"jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
|
|
)
|
|
return False
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def require_jupyter_kernel_gateway_installed() -> Callable[[T], T]:
|
|
"""Decorator that checks if jupyter-kernel-gateway is installed before function execution.
|
|
|
|
Returns:
|
|
Callable[[T], T]: A decorator function that either:
|
|
- Returns the original function unchanged if jupyter-kernel-gateway is installed
|
|
- Returns a patched version of the function that will raise a helpful error indicating the missing dependency when called
|
|
"""
|
|
if is_jupyter_kernel_gateway_installed():
|
|
|
|
def decorator(o: T) -> T:
|
|
return o
|
|
|
|
else:
|
|
|
|
def decorator(o: T) -> T:
|
|
return patch_object(o, missing_modules={}, dep_target="jupyter-executor")
|
|
|
|
return decorator
|
|
|
|
|
|
def skip_on_missing_jupyter_kernel_gateway() -> Callable[[T], T]:
|
|
"""Decorator to skip a test if an optional module is missing"""
|
|
# Add pytest.mark.jupyter_executor decorator
|
|
mark_name = "jupyter_executor"
|
|
|
|
if is_jupyter_kernel_gateway_installed():
|
|
|
|
def decorator(o: T) -> T:
|
|
import pytest
|
|
|
|
pytest_mark_o = getattr(pytest.mark, mark_name)(o)
|
|
return pytest_mark_o # type: ignore[no-any-return]
|
|
|
|
else:
|
|
|
|
def decorator(o: T) -> T:
|
|
import pytest
|
|
|
|
pytest_mark_o = getattr(pytest.mark, mark_name)(o)
|
|
return pytest.mark.skip( # type: ignore[return-value,no-any-return]
|
|
reason="jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
|
|
)(pytest_mark_o)
|
|
|
|
return decorator
|