snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/_internal/telemetry.py

@@ -8,25 +8,13 @@ import sys
 import time
 import traceback
 import types
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Mapping,
-    Optional,
-    Set,
-    Tuple,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import Any, Callable, Iterable, Mapping, Optional, TypeVar, Union, cast
 
 from typing_extensions import ParamSpec
 
 from snowflake import connector
 from snowflake.connector import telemetry as connector_telemetry, time_util
+from snowflake.ml import version as snowml_version
 from snowflake.ml._internal import env
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -99,13 +87,13 @@ class _TelemetrySourceType(enum.Enum):
     AUGMENT_TELEMETRY = "SNOWML_AUGMENT_TELEMETRY"
 
 
-_statement_params_context_var: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar("statement_params")
+_statement_params_context_var: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("statement_params")
 
 
 class _StatementParamsPatchManager:
     def __init__(self) -> None:
-        self._patch_cache: Set[server_connection.ServerConnection] = set()
-        self._context_var: contextvars.ContextVar[Dict[str, str]] = _statement_params_context_var
+        self._patch_cache: set[server_connection.ServerConnection] = set()
+        self._context_var: contextvars.ContextVar[dict[str, str]] = _statement_params_context_var
 
     def apply_patches(self) -> None:
         try:
@@ -117,7 +105,7 @@ class _StatementParamsPatchManager:
         except snowpark_exceptions.SnowparkSessionException:
             pass
 
-    def set_statement_params(self, statement_params: Dict[str, str]) -> None:
+    def set_statement_params(self, statement_params: dict[str, str]) -> None:
         # Only set value if not already set in context
         if not self._context_var.get({}):
             self._context_var.set(statement_params)
@@ -152,7 +140,6 @@ class _StatementParamsPatchManager:
                 if throw_on_patch_fail:  # primarily used for testing
                     raise
                 # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
-                pass
 
     def _patch_with_statement_params(
         self, target: object, function_name: str, param_name: str = "statement_params"
@@ -197,10 +184,10 @@ class _StatementParamsPatchManager:
 
         setattr(target, function_name, wrapper)
 
-    def __getstate__(self) -> Dict[str, Any]:
+    def __getstate__(self) -> dict[str, Any]:
         return {}
 
-    def __setstate__(self, state: Dict[str, Any]) -> None:
+    def __setstate__(self, state: dict[str, Any]) -> None:
         # unpickling does not call __init__ by default, do it manually here
         self.__init__()  # type: ignore[misc]
 
@@ -210,7 +197,7 @@ _patch_manager = _StatementParamsPatchManager()
 
 def get_statement_params(
     project: str, subproject: Optional[str] = None, class_name: Optional[str] = None
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Get telemetry statement parameters.
 
@@ -231,8 +218,8 @@ def get_statement_params(
 
 
 def add_statement_params_custom_tags(
-    statement_params: Optional[Dict[str, Any]], custom_tags: Mapping[str, Any]
-) -> Dict[str, Any]:
+    statement_params: Optional[dict[str, Any]], custom_tags: Mapping[str, Any]
+) -> dict[str, Any]:
     """
     Add custom_tags to existing statement_params.  Overwrite keys in custom_tags dict that already exist.
     If existing statement_params are not provided, do nothing as the information cannot be effectively tracked.
@@ -246,7 +233,7 @@ def add_statement_params_custom_tags(
     """
     if not statement_params:
         return {}
-    existing_custom_tags: Dict[str, Any] = statement_params.pop(TelemetryField.KEY_CUSTOM_TAGS.value, {})
+    existing_custom_tags: dict[str, Any] = statement_params.pop(TelemetryField.KEY_CUSTOM_TAGS.value, {})
     existing_custom_tags.update(custom_tags)
     # NOTE: This can be done with | operator after upgrade from py3.8
     return {
@@ -289,17 +276,17 @@ def get_function_usage_statement_params(
     *,
     function_category: str = TelemetryField.FUNC_CAT_USAGE.value,
     function_name: Optional[str] = None,
-    function_parameters: Optional[Dict[str, Any]] = None,
+    function_parameters: Optional[dict[str, Any]] = None,
     api_calls: Optional[
-        List[
+        list[
             Union[
-                Dict[str, Union[Callable[..., Any], str]],
+                dict[str, Union[Callable[..., Any], str]],
                 Union[Callable[..., Any], str],
             ]
         ]
     ] = None,
-    custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
-) -> Dict[str, Any]:
+    custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
+) -> dict[str, Any]:
     """
     Get function usage statement parameters.
 
@@ -321,12 +308,12 @@ def get_function_usage_statement_params(
         >>> df.collect(statement_params=statement_params)
     """
     telemetry_type = f"{env.SOURCE.lower()}_{TelemetryField.TYPE_FUNCTION_USAGE.value}"
-    statement_params: Dict[str, Any] = {
+    statement_params: dict[str, Any] = {
         connector_telemetry.TelemetryField.KEY_SOURCE.value: env.SOURCE,
         TelemetryField.KEY_PROJECT.value: project,
         TelemetryField.KEY_SUBPROJECT.value: subproject,
         TelemetryField.KEY_OS.value: env.OS,
-        TelemetryField.KEY_VERSION.value: env.VERSION,
+        TelemetryField.KEY_VERSION.value: snowml_version.VERSION,
         TelemetryField.KEY_PYTHON_VERSION.value: env.PYTHON_VERSION,
         connector_telemetry.TelemetryField.KEY_TYPE.value: telemetry_type,
         TelemetryField.KEY_CATEGORY.value: function_category,
@@ -339,7 +326,7 @@ def get_function_usage_statement_params(
     if api_calls:
         statement_params[TelemetryField.KEY_API_CALLS.value] = []
         for api_call in api_calls:
-            if isinstance(api_call, Dict):
+            if isinstance(api_call, dict):
                 telemetry_api_call = api_call.copy()
                 # convert Callable to str
                 for field, api in api_call.items():
@@ -388,7 +375,7 @@ def send_custom_usage(
     *,
     telemetry_type: str,
     subproject: Optional[str] = None,
-    data: Optional[Dict[str, Any]] = None,
+    data: Optional[dict[str, Any]] = None,
     **kwargs: Any,
 ) -> None:
     active_session = next(iter(session._get_active_sessions()))
@@ -409,17 +396,17 @@ def send_api_usage_telemetry(
     api_calls_extractor: Optional[
         Callable[
             ...,
-            List[
+            list[
                 Union[
-                    Dict[str, Union[Callable[..., Any], str]],
+                    dict[str, Union[Callable[..., Any], str]],
                     Union[Callable[..., Any], str],
                 ]
             ],
         ]
     ] = None,
-    sfqids_extractor: Optional[Callable[..., List[str]]] = None,
+    sfqids_extractor: Optional[Callable[..., list[str]]] = None,
     subproject_extractor: Optional[Callable[[Any], str]] = None,
-    custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
+    custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
     """
     Decorator that sends API usage telemetry and adds function usage statement parameters to the dataframe returned by
@@ -454,7 +441,7 @@ def send_api_usage_telemetry(
         def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
             params = _get_func_params(func, func_params_to_log, args, kwargs) if func_params_to_log else None
 
-            api_calls: List[Union[Dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []
+            api_calls: list[Union[dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []
             if api_calls_extractor:
                 extracted_api_calls = api_calls_extractor(args[0])
                 for api_call in extracted_api_calls:
@@ -484,7 +471,7 @@ def send_api_usage_telemetry(
                 custom_tags=custom_tags,
             )
 
-            def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: Dict[str, Any]) -> _ReturnValue:
+            def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: dict[str, Any]) -> _ReturnValue:
                 """
                 Update SnowML function usage statement parameters to the object if it is a Snowpark DataFrame.
                 Used to track APIs returning a Snowpark DataFrame.
@@ -614,7 +601,7 @@ def _get_full_func_name(func: Callable[..., Any]) -> str:
 
 def _get_func_params(
     func: Callable[..., Any], func_params_to_log: Optional[Iterable[str]], args: Any, kwargs: Any
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Get function parameters.
 
@@ -639,7 +626,7 @@ def _get_func_params(
     return params
 
 
-def _extract_arg_value(field: str, func_spec: inspect.FullArgSpec, args: Any, kwargs: Any) -> Tuple[bool, Any]:
+def _extract_arg_value(field: str, func_spec: inspect.FullArgSpec, args: Any, kwargs: Any) -> tuple[bool, Any]:
     """
     Function to extract a specified argument value.
 
@@ -702,11 +689,11 @@ class _SourceTelemetryClient:
         self.source: str = env.SOURCE
         self.project: Optional[str] = project
         self.subproject: Optional[str] = subproject
-        self.version = env.VERSION
+        self.version = snowml_version.VERSION
         self.python_version: str = env.PYTHON_VERSION
         self.os: str = env.OS
 
-    def _send(self, msg: Dict[str, Any], timestamp: Optional[int] = None) -> None:
+    def _send(self, msg: dict[str, Any], timestamp: Optional[int] = None) -> None:
         """
         Add telemetry data to a batch in connector client.
 
@@ -720,7 +707,7 @@ class _SourceTelemetryClient:
             telemetry_data = connector_telemetry.TelemetryData(message=msg, timestamp=timestamp)
             self._telemetry.try_add_log_to_batch(telemetry_data)
 
-    def _create_basic_telemetry_data(self, telemetry_type: str) -> Dict[str, Any]:
+    def _create_basic_telemetry_data(self, telemetry_type: str) -> dict[str, Any]:
         message = {
             connector_telemetry.TelemetryField.KEY_SOURCE.value: self.source,
             TelemetryField.KEY_PROJECT.value: self.project,
@@ -738,10 +725,10 @@ class _SourceTelemetryClient:
         func_name: str,
         function_category: str,
         duration: float,
-        func_params: Optional[Dict[str, Any]] = None,
-        api_calls: Optional[List[Dict[str, Any]]] = None,
-        sfqids: Optional[List[Any]] = None,
-        custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
+        func_params: Optional[dict[str, Any]] = None,
+        api_calls: Optional[list[dict[str, Any]]] = None,
+        sfqids: Optional[list[Any]] = None,
+        custom_tags: Optional[dict[str, Union[bool, int, str, float]]] = None,
         error: Optional[str] = None,
         error_code: Optional[str] = None,
         stack_trace: Optional[str] = None,
@@ -761,7 +748,7 @@ class _SourceTelemetryClient:
             error_code: Error code.
             stack_trace: Error stack trace.
         """
-        data: Dict[str, Any] = {
+        data: dict[str, Any] = {
             TelemetryField.KEY_FUNC_NAME.value: func_name,
             TelemetryField.KEY_CATEGORY.value: function_category,
         }
@@ -775,7 +762,7 @@ class _SourceTelemetryClient:
             data[TelemetryField.KEY_CUSTOM_TAGS.value] = custom_tags
 
         telemetry_type = f"{self.source.lower()}_{TelemetryField.TYPE_FUNCTION_USAGE.value}"
-        message: Dict[str, Any] = {
+        message: dict[str, Any] = {
             **self._create_basic_telemetry_data(telemetry_type),
             TelemetryField.KEY_DATA.value: data,
             TelemetryField.KEY_DURATION.value: duration,
@@ -795,7 +782,7 @@ class _SourceTelemetryClient:
             self._telemetry.send_batch()
 
 
-def get_sproc_statement_params_kwargs(sproc: Callable[..., Any], statement_params: Dict[str, Any]) -> Dict[str, Any]:
+def get_sproc_statement_params_kwargs(sproc: Callable[..., Any], statement_params: dict[str, Any]) -> dict[str, Any]:
     """
     Get statement_params keyword argument for sproc call.
 

snowflake/ml/_internal/type_utils.py

@@ -11,7 +11,7 @@ T = TypeVar("T")
 class LazyType(Generic[T]):
     """Utility type to help defer need of importing."""
 
-    def __init__(self, klass: Union[str, Type[T]]) -> None:
+    def __init__(self, klass: Union[str, type[T]]) -> None:
         self.qualname = ""
         if isinstance(klass, str):
             parts = klass.rsplit(".", 1)
@@ -30,7 +30,7 @@ class LazyType(Generic[T]):
         return self.isinstance(obj)
 
     @classmethod
-    def from_type(cls, typ_: Union["LazyType[T]", Type[T]]) -> "LazyType[T]":
+    def from_type(cls, typ_: Union["LazyType[T]", type[T]]) -> "LazyType[T]":
         if isinstance(typ_, LazyType):
             return typ_
         return cls(typ_)
@@ -48,7 +48,7 @@ class LazyType(Generic[T]):
     def __repr__(self) -> str:
         return f'LazyType("{self.module}", "{self.qualname}")'
 
-    def get_class(self) -> Type[T]:
+    def get_class(self) -> type[T]:
         if self._runtime_class is None:
             try:
                 m = importlib.import_module(self.module)

snowflake/ml/_internal/utils/db_utils.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.snowpark import session
@@ -19,7 +19,7 @@ def db_object_exists(
     *,
     database_name: Optional[sql_identifier.SqlIdentifier] = None,
     schema_name: Optional[sql_identifier.SqlIdentifier] = None,
-    statement_params: Optional[Dict[str, Any]] = None,
+    statement_params: Optional[dict[str, Any]] = None,
 ) -> bool:
     """Check if object exists in database.
 

snowflake/ml/_internal/utils/identifier.py

@@ -1,5 +1,5 @@
 import re
-from typing import Any, List, Optional, Tuple, Union, overload
+from typing import Any, Optional, Union, overload
 
 from snowflake.snowpark._internal.analyzer import analyzer_utils
 
@@ -112,7 +112,7 @@ def get_inferred_name(name: str) -> str:
     return escaped_id
 
 
-def concat_names(names: List[str]) -> str:
+def concat_names(names: list[str]) -> str:
     """Concatenates `names` to form one valid id.
 
 
@@ -142,7 +142,7 @@ def rename_to_valid_snowflake_identifier(name: str) -> str:
 
 def parse_schema_level_object_identifier(
     object_name: str,
-) -> Tuple[Union[str, Any], Union[str, Any], Union[str, Any]]:
+) -> tuple[Union[str, Any], Union[str, Any], Union[str, Any]]:
     """Parse a string which starts with schema level object.
 
     Args:
@@ -172,7 +172,7 @@ def parse_schema_level_object_identifier(
 
 def parse_snowflake_stage_path(
     path: str,
-) -> Tuple[Union[str, Any], Union[str, Any], Union[str, Any], Union[str, Any]]:
+) -> tuple[Union[str, Any], Union[str, Any], Union[str, Any], Union[str, Any]]:
     """Parse a string which represents a snowflake stage path.
 
     Args:
@@ -260,11 +260,11 @@ def get_unescaped_names(ids: str) -> str:
 
 
 @overload
-def get_unescaped_names(ids: List[str]) -> List[str]:
+def get_unescaped_names(ids: list[str]) -> list[str]:
     ...
 
 
-def get_unescaped_names(ids: Optional[Union[str, List[str]]]) -> Optional[Union[str, List[str]]]:
+def get_unescaped_names(ids: Optional[Union[str, list[str]]]) -> Optional[Union[str, list[str]]]:
     """Given a user provided identifier(s), this method will compute the equivalent column name identifier(s) in the
     response pandas dataframe(i.e., in the response of snowpark_df.to_pandas()) using the rules defined here
     https://docs.snowflake.com/en/sql-reference/identifiers-syntax.
@@ -308,11 +308,11 @@ def get_inferred_names(names: str) -> str:
 
 
 @overload
-def get_inferred_names(names: List[str]) -> List[str]:
+def get_inferred_names(names: list[str]) -> list[str]:
     ...
 
 
-def get_inferred_names(names: Optional[Union[str, List[str]]]) -> Optional[Union[str, List[str]]]:
+def get_inferred_names(names: Optional[Union[str, list[str]]]) -> Optional[Union[str, list[str]]]:
     """Given a user provided *string(s)*, this method will compute the equivalent column name identifier(s)
     in case of column name contains special characters, and maintains case-sensitivity
     https://docs.snowflake.com/en/sql-reference/identifiers-syntax.

snowflake/ml/_internal/utils/import_utils.py

@@ -1,5 +1,5 @@
 import importlib
-from typing import Any, Tuple
+from typing import Any
 
 
 class MissingOptionalDependency:
@@ -46,7 +46,7 @@ def import_with_fallbacks(*targets: str) -> Any:
     raise ImportError(f"None of the requested targets could be imported. Requested: {', '.join(targets)}")
 
 
-def import_or_get_dummy(target: str) -> Tuple[Any, bool]:
+def import_or_get_dummy(target: str) -> tuple[Any, bool]:
     """Try to import the the given target or return a dummy object.
 
     If the import target (package/module/symbol) is available, the target will be returned. If it is not available,

snowflake/ml/_internal/utils/parallelize.py

@@ -1,7 +1,7 @@
 import math
 from contextlib import contextmanager
 from timeit import default_timer
-from typing import Any, Callable, Dict, Generator, Iterable, List, Optional
+from typing import Any, Callable, Generator, Iterable, Optional
 
 import snowflake.snowpark.functions as F
 from snowflake import snowpark
@@ -17,17 +17,17 @@ def timer() -> Generator[Callable[[], float], None, None]:
     yield lambda: elapser()
 
 
-def _flatten(L: Iterable[List[Any]]) -> List[Any]:
+def _flatten(L: Iterable[list[Any]]) -> list[Any]:
     return [val for sublist in L for val in sublist]
 
 
 def map_dataframe_by_column(
     df: snowpark.DataFrame,
-    cols: List[str],
-    map_func: Callable[[snowpark.DataFrame, List[str]], snowpark.DataFrame],
+    cols: list[str],
+    map_func: Callable[[snowpark.DataFrame, list[str]], snowpark.DataFrame],
     partition_size: int,
-    statement_params: Optional[Dict[str, Any]] = None,
-) -> List[List[Any]]:
+    statement_params: Optional[dict[str, Any]] = None,
+) -> list[list[Any]]:
     """Applies the `map_func` to the input DataFrame by parallelizing it over subsets of the column.
 
     Because the return results are materialized as Python lists *in memory*, this method should
@@ -84,7 +84,7 @@ def map_dataframe_by_column(
             unioned_df = mapped_df if unioned_df is None else unioned_df.union_all(mapped_df)
 
     # Store results in a list of size |n_partitions| x |n_rows| x |n_output_cols|
-    all_results: List[List[List[Any]]] = [[] for _ in range(n_partitions - 1)]
+    all_results: list[list[list[Any]]] = [[] for _ in range(n_partitions - 1)]
 
     # Collect the results of the first n-1 partitions, removing the partition_id column
     unioned_result = unioned_df.collect(statement_params=statement_params) if unioned_df is not None else []

snowflake/ml/_internal/utils/pkg_version_utils.py

@@ -1,6 +1,6 @@
 import sys
 import warnings
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Optional, Union
 
 from packaging.version import Version
 
@@ -8,7 +8,7 @@ from snowflake.ml._internal import telemetry
 from snowflake.snowpark import AsyncJob, Row, Session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
-cache: Dict[str, Optional[str]] = {}
+cache: dict[str, Optional[str]] = {}
 
 _PROJECT = "ModelDevelopment"
 _SUBPROJECT = "utils"
@@ -23,8 +23,8 @@ def is_relaxed() -> bool:
 
 
 def get_valid_pkg_versions_supported_in_snowflake_conda_channel(
-    pkg_versions: List[str], session: Session, subproject: Optional[str] = None
-) -> List[str]:
+    pkg_versions: list[str], session: Session, subproject: Optional[str] = None
+) -> list[str]:
     if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
         return pkg_versions
     else:
@@ -32,9 +32,9 @@ def get_valid_pkg_versions_supported_in_snowflake_conda_channel(
 
 
 def _get_valid_pkg_versions_supported_in_snowflake_conda_channel_async(
-    pkg_versions: List[str], session: Session, subproject: Optional[str] = None
-) -> List[str]:
-    pkg_version_async_job_list: List[Tuple[str, AsyncJob]] = []
+    pkg_versions: list[str], session: Session, subproject: Optional[str] = None
+) -> list[str]:
+    pkg_version_async_job_list: list[tuple[str, AsyncJob]] = []
     for pkg_version in pkg_versions:
         if pkg_version not in cache:
             # Execute pkg version queries asynchronously.
@@ -64,7 +64,7 @@ def _get_valid_pkg_versions_supported_in_snowflake_conda_channel_async(
 
 def _query_pkg_version_supported_in_snowflake_conda_channel(
     pkg_version: str, session: Session, block: bool, subproject: Optional[str] = None
-) -> Union[AsyncJob, List[Row]]:
+) -> Union[AsyncJob, list[Row]]:
     tokens = pkg_version.split("==")
     if len(tokens) != 2:
         raise RuntimeError(
@@ -102,9 +102,9 @@ def _query_pkg_version_supported_in_snowflake_conda_channel(
     return pkg_version_list_or_async_job
 
 
-def _get_conda_packages_and_emit_warnings(pkg_versions: List[str]) -> List[str]:
-    pkg_version_conda_list: List[str] = []
-    pkg_version_warning_list: List[List[str]] = []
+def _get_conda_packages_and_emit_warnings(pkg_versions: list[str]) -> list[str]:
+    pkg_version_conda_list: list[str] = []
+    pkg_version_warning_list: list[list[str]] = []
     for pkg_version in pkg_versions:
         try:
             conda_pkg_version = cache[pkg_version]

snowflake/ml/_internal/utils/query_result_checker.py

@@ -1,7 +1,7 @@
 from __future__ import annotations  # for return self methods
 
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Optional
 
 from snowflake import connector, snowpark
 from snowflake.ml._internal.utils import formatting
@@ -123,7 +123,7 @@ def cell_value_by_column_matcher(
     return True
 
 
-_DEFAULT_MATCHERS: List[Callable[[List[snowpark.Row], Optional[str]], bool]] = [
+_DEFAULT_MATCHERS: list[Callable[[list[snowpark.Row], Optional[str]], bool]] = [
     partial(result_dimension_matcher, 1, 1),
     partial(column_name_matcher, "status"),
 ]
@@ -252,12 +252,12 @@ class SqlResultValidator(ResultValidator):
     """
 
     def __init__(
-        self, session: snowpark.Session, query: str, statement_params: Optional[Dict[str, Any]] = None
+        self, session: snowpark.Session, query: str, statement_params: Optional[dict[str, Any]] = None
     ) -> None:
         self._session: snowpark.Session = session
         self._query: str = query
         self._success_matchers: list[Callable[[list[snowpark.Row], Optional[str]], bool]] = []
-        self._statement_params: Optional[Dict[str, Any]] = statement_params
+        self._statement_params: Optional[dict[str, Any]] = statement_params
 
     def _get_result(self) -> list[snowpark.Row]:
         """Collect the result of the given SQL query."""

snowflake/ml/_internal/utils/snowflake_env.py

@@ -1,15 +1,15 @@
 import enum
-from typing import Any, Dict, Optional, TypedDict, cast
+from typing import Any, Optional, TypedDict, cast
 
 from packaging import version
 from typing_extensions import NotRequired, Required
 
 from snowflake.ml._internal.utils import query_result_checker
-from snowflake.snowpark import session
+from snowflake.snowpark import exceptions as sp_exceptions, session
 
 
 def get_current_snowflake_version(
-    sess: session.Session, *, statement_params: Optional[Dict[str, Any]] = None
+    sess: session.Session, *, statement_params: Optional[dict[str, Any]] = None
 ) -> version.Version:
     """Get Snowflake Version as a version.Version object follow PEP way of versioning, that is to say:
         "7.44.2 b202312132139364eb71238" to <Version('7.44.2+b202312132139364eb71238')>
@@ -60,8 +60,8 @@ class SnowflakeRegion(TypedDict):
 
 
 def get_regions(
-    sess: session.Session, *, statement_params: Optional[Dict[str, Any]] = None
-) -> Dict[str, SnowflakeRegion]:
+    sess: session.Session, *, statement_params: Optional[dict[str, Any]] = None
+) -> dict[str, SnowflakeRegion]:
     res = (
         query_result_checker.SqlResultValidator(sess, "SHOW REGIONS", statement_params=statement_params)
         .has_column("snowflake_region")
@@ -93,7 +93,7 @@ def get_regions(
     return res_dict
 
 
-def get_current_region_id(sess: session.Session, *, statement_params: Optional[Dict[str, Any]] = None) -> str:
+def get_current_region_id(sess: session.Session, *, statement_params: Optional[dict[str, Any]] = None) -> str:
     res = (
         query_result_checker.SqlResultValidator(
             sess, "SELECT CURRENT_REGION() AS CURRENT_REGION", statement_params=statement_params
@@ -103,3 +103,25 @@ def get_current_region_id(sess: session.Session, *, statement_params: Optional[D
     )
 
     return cast(str, res.CURRENT_REGION)
+
+
+def get_current_cloud(
+    sess: session.Session,
+    default: Optional[SnowflakeCloudType] = None,
+    *,
+    statement_params: Optional[dict[str, Any]] = None,
+) -> SnowflakeCloudType:
+    region_id = get_current_region_id(sess, statement_params=statement_params)
+    try:
+        region = get_regions(sess, statement_params=statement_params)[region_id]
+        return region["cloud"]
+    except sp_exceptions.SnowparkSQLException:
+        # SHOW REGIONS not available, try to infer cloud from region name
+        region_name = region_id.split(".", 1)[-1]  # Drop region group if any, e.g. PUBLIC
+        cloud_name_maybe = region_name.split("_", 1)[0]  # Extract cloud name, e.g. AWS_US_WEST -> AWS
+        try:
+            return SnowflakeCloudType.from_value(cloud_name_maybe)
+        except ValueError:
+            if default:
+                return default
+            raise

snowflake/ml/_internal/utils/snowpark_dataframe_utils.py

@@ -1,13 +1,13 @@
 import logging
 import warnings
-from typing import List, Optional
+from typing import Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.snowpark import functions, types
 
 
-def cast_snowpark_dataframe(df: snowpark.DataFrame, ignore_columns: Optional[List[str]] = None) -> snowpark.DataFrame:
+def cast_snowpark_dataframe(df: snowpark.DataFrame, ignore_columns: Optional[list[str]] = None) -> snowpark.DataFrame:
     """Cast columns in the dataframe to types that are compatible with tensor.
 
     It assists FileSet.make() in performing implicit data casting.

snowflake/ml/_internal/utils/sql_identifier.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 from snowflake.ml._internal.utils import identifier
 
@@ -77,13 +77,13 @@ class SqlIdentifier(str):
         return super().__hash__()
 
 
-def to_sql_identifiers(list_of_str: List[str], *, case_sensitive: bool = False) -> List[SqlIdentifier]:
+def to_sql_identifiers(list_of_str: list[str], *, case_sensitive: bool = False) -> list[SqlIdentifier]:
     return [SqlIdentifier(val, case_sensitive=case_sensitive) for val in list_of_str]
 
 
 def parse_fully_qualified_name(
     name: str,
-) -> Tuple[Optional[SqlIdentifier], Optional[SqlIdentifier], SqlIdentifier]:
+) -> tuple[Optional[SqlIdentifier], Optional[SqlIdentifier], SqlIdentifier]:
     db, schema, object = identifier.parse_schema_level_object_identifier(name)
 
     assert name is not None, f"Unable parse the input name `{name}` as fully qualified."