Repository: apache/airflow
File: providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py (Python)

Code before:
        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
        or you can pass a string as "v1.0" or "beta".
    """

    DEFAULT_HEADERS = {"Accept": "application/json;q=1"}
    DEFAULT_SCOPE = "https://graph.microsoft.com/.default"
    cached_request_adapters: dict[str, tuple[str, RequestAdapter]] = {}
    conn_type: str = "msgraph"
    conn_name_attr: str = "conn_id"
    default_conn_name: str = "msgraph_default"
    hook_name: str = "Microsoft Graph API"

    def __init__(
        self,
        conn_id: str = default_conn_name,
        timeout: float | None = None,
        proxies: dict | None = None,
        host: str | None = None,
        scopes: str | list[str] | None = None,
        api_version: APIVersion | str | None = None,
    ):
        super().__init__()
        self.conn_id = conn_id
        self.timeout = timeout
        self.proxies = proxies
        self.host = self._ensure_protocol(host) if host else None
        if isinstance(scopes, str):
            self.scopes = [scopes]
        else:
            self.scopes = scopes or [self.DEFAULT_SCOPE]
        self.api_version = self.resolve_api_version_from_value(api_version)

    def _ensure_protocol(self, url: str) -> str:
        if not url:
            return url
        if url.startswith(("http://", "https://")):
            return url

        self.log.warning(
            "URL '%s' is missing protocol prefix. Automatically adding 'https://'. "
            "Please update your connection configuration to include the full URL with protocol.",
            url,
        )
        return f"https://{url}"

    @classmethod
    def get_connection_form_widgets(cls) -> dict[str, Any]:
        """Return connection widgets to add to connection form."""
        from flask_appbuilder.fieldwidgets import BS3TextAreaFieldWidget, BS3TextFieldWidget
        from flask_babel import lazy_gettext
        from wtforms.fields import BooleanField, StringField

Review comment (refactor):
Since _ensure_protocol already checks for if not url, can we remove the if host else None check here to keep it DRY?

Code after:
        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
        or you can pass a string as "v1.0" or "beta".
    """

    DEFAULT_HEADERS = {"Accept": "application/json;q=1"}
    DEFAULT_SCOPE = "https://graph.microsoft.com/.default"
    cached_request_adapters: dict[str, tuple[str, RequestAdapter]] = {}
    conn_type: str = "msgraph"
    conn_name_attr: str = "conn_id"
    default_conn_name: str = "msgraph_default"
    hook_name: str = "Microsoft Graph API"

    def __init__(
        self,
        conn_id: str = default_conn_name,
        timeout: float | None = None,
        proxies: dict | None = None,
        host: str | None = None,
        scopes: str | list[str] | None = None,
        api_version: APIVersion | str | None = None,
    ):
        super().__init__()
        self.conn_id = conn_id
        self.timeout = timeout
        self.proxies = proxies
        self.host = self._ensure_protocol(host)
        if isinstance(scopes, str):
            self.scopes = [scopes]
        else:
            self.scopes = scopes or [self.DEFAULT_SCOPE]
        self.api_version = self.resolve_api_version_from_value(api_version)

    def _ensure_protocol(self, host: str | None, schema: str = "https") -> str | None:
        """Ensure URL has http:// or https:// protocol prefix."""
        if not host:
            return None

        if host.startswith(("http://", "https://")):
            return host

        self.log.warning(
            "URL '%s' is missing protocol prefix. Automatically adding '%s://'. "
            "Please update your connection configuration to include the full URL with protocol.",
            host,
            schema,
        )
        return f"{schema}://{host}"

    @classmethod
    def get_connection_form_widgets(cls) -> dict[str, Any]:
        """Return connection widgets to add to connection form."""