Repository: ray-project/ray
File: python/ray/serve/taskiq_task_processor.py (Python)

Code before:
        "import": "taskiq_aio_kafka.AioKafkaBroker",
        "package": "taskiq-aio-kafka",
        "queue_param": "kafka_topic",
        "required_kwargs": ["bootstrap_servers"],
    },
}


def _import_broker_class(broker_type: str):
    """Lazily import and return the broker class for the given broker_type."""
    entry = _BROKER_REGISTRY.get(broker_type)
    if entry is None:
        raise ValueError(
            f"Unsupported broker_type: {broker_type!r}. "
            f"Supported types: {list(_BROKER_REGISTRY.keys())}"
        )

    module_path, class_name = entry["import"].rsplit(".", 1)
    try:
        module = importlib.import_module(module_path)
        return getattr(module, class_name)
    except ImportError as e:
        raise ImportError(
            f"Broker {broker_type!r} requires package {entry['package']!r}. "
            f"Install it with: pip install {entry['package']}"
        ) from e


def _create_broker(
    broker_type: str,
    queue_name: str,
    broker_kwargs: Optional[Dict[str, Any]] = None,
):
    """Create a broker instance from the given config."""
    broker_cls = _import_broker_class(broker_type)
    entry = _BROKER_REGISTRY[broker_type]

    kwargs = dict(broker_kwargs) if broker_kwargs else {}

    # Validate required kwargs are present.
    required = entry.get("required_kwargs", [])
    missing = [k for k in required if k not in kwargs]
    if missing:
        raise ValueError(
            f"Broker {broker_type!r} requires the following keys in "
            f"broker_kwargs: {missing}"
        )

    # Inject the queue/topic name under the broker-specific parameter name.
    queue_param = entry["queue_param"]
    kwargs[queue_param] = queue_name

Review comment (refactor):
could use this utility https://github.com/ray-project/ray/blob/master/python/ray/_common/utils.py#L60

Code after:
        "import": "taskiq_aio_kafka.AioKafkaBroker",
        "package": "taskiq-aio-kafka",
        "queue_param": "kafka_topic",
        "required_kwargs": ["bootstrap_servers"],
    },
}


def _import_broker_class(broker_type: str):
    """Lazily import and return the broker class for the given broker_type."""
    entry = _BROKER_REGISTRY.get(broker_type)
    if entry is None:
        raise ValueError(
            f"Unsupported broker_type: {broker_type!r}. "
            f"Supported types: {list(_BROKER_REGISTRY.keys())}"
        )

    try:
        return import_attr(entry["import"])
    except (ImportError, ModuleNotFoundError) as e:
        raise ImportError(
            f"Broker {broker_type!r} requires package {entry['package']!r}. "
            f"Install it with: pip install {entry['package']}"
        ) from e


def _create_broker(
    broker_type: str,
    queue_name: str,
    broker_kwargs: Optional[Dict[str, Any]] = None,
):
    """Create a broker instance from the given config."""
    entry = _BROKER_REGISTRY.get(broker_type)
    if entry is None:
        raise ValueError(
            f"Unsupported broker_type: {broker_type!r}. "
            f"Supported types: {list(_BROKER_REGISTRY.keys())}"
        )

    kwargs = dict(broker_kwargs) if broker_kwargs else {}

    # Validate required kwargs are present.
    required = entry.get("required_kwargs", [])
    missing = [k for k in required if k not in kwargs]
    if missing:
        raise ValueError(
            f"Broker {broker_type!r} requires the following keys in "
            f"broker_kwargs: {missing}"
        )

    broker_cls = _import_broker_class(broker_type)