maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/lib/wrapped_pickle.py

@@ -0,0 +1,139 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextvars
+import functools
+import io
+import os
+import pickle as raw_pickle
+import sys
+from typing import Callable
+
+import cloudpickle
+
+if sys.version_info[:2] < (3, 8):  # pragma: no cover
+    try:
+        import pickle5 as pickle_mod  # nosec  # pylint: disable=import_pickle
+    except ImportError:
+        import pickle as pickle_mod  # nosec  # pylint: disable=import_pickle
+else:
+    import pickle as pickle_mod  # nosec  # pylint: disable=import_pickle
+
+__all__ = [
+    "Pickler",
+    "Unpickler",
+    "dump",
+    "dumps",
+    "load",
+    "loads",
+    "PickleError",
+    "PicklingError",
+    "UnpicklingError",
+    "DEFAULT_PROTOCOL",
+    "HIGHEST_PROTOCOL",
+    "switch_unpickle",
+]
+
+Pickler = cloudpickle.Pickler
+dump = cloudpickle.dump
+dumps = cloudpickle.dumps
+
+PickleError = pickle_mod.PickleError
+PicklingError = pickle_mod.PicklingError
+UnpicklingError = pickle_mod.UnpicklingError
+
+DEFAULT_PROTOCOL = pickle_mod.DEFAULT_PROTOCOL
+HIGHEST_PROTOCOL = pickle_mod.HIGHEST_PROTOCOL
+
+
+_default_forbidden_val = bool(int(os.getenv("MAXFRAME_FORBIDDEN_PICKLE", "0")))
+_unpickle_forbidden_var = contextvars.ContextVar(
+    "unpickle_forbidden", default=_default_forbidden_val
+)
+
+
+def is_unpickle_forbidden():
+    return _unpickle_forbidden_var.get(False)
+
+
+class Unpickler(pickle_mod.Unpickler):
+    @functools.wraps(pickle_mod.Unpickler.load)
+    def load(self):
+        if _unpickle_forbidden_var.get(False):
+            raise ValueError("Unpickle is forbidden here")
+        return super().load()
+
+
+@functools.wraps(pickle_mod.load)
+def load(file, **kwargs):
+    return Unpickler(file, **kwargs).load()
+
+
+@functools.wraps(pickle_mod.loads)
+def loads(s, **kwargs):
+    if isinstance(s, str):  # pragma: no cover
+        raise TypeError("Can't load pickle from unicode string")
+    file = io.BytesIO(s)
+    return Unpickler(file, **kwargs).load()
+
+
+# patch original pickle methods
+pickle_mod.Unpickler = Unpickler
+pickle_mod.load = load
+pickle_mod.loads = loads
+
+raw_pickle.Unpickler = Unpickler
+raw_pickle.load = load
+raw_pickle.loads = loads
+
+cloudpickle.Unpickler = Unpickler
+cloudpickle.load = load
+cloudpickle.loads = loads
+
+
+class _UnpickleSwitch:
+    def __init__(self, forbidden: bool = True):
+        self._token = None
+        self._forbidden = forbidden
+
+    def __enter__(self):
+        self._token = _unpickle_forbidden_var.set(self._forbidden)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        _unpickle_forbidden_var.reset(self._token)
+
+    def __call__(self, func: Callable) -> Callable:
+        if asyncio.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def wrapped(*args, **kwargs):
+                with _UnpickleSwitch(forbidden=self._forbidden):
+                    return await func(*args, **kwargs)
+
+        else:
+
+            @functools.wraps(func)
+            def wrapped(*args, **kwargs):
+                with _UnpickleSwitch(forbidden=self._forbidden):
+                    return func(*args, **kwargs)
+
+        return wrapped
+
+
+def switch_unpickle(func=None, *, forbidden: bool = True):
+    switch_obj = _UnpickleSwitch(forbidden=forbidden)
+    if func is None:
+        return switch_obj
+    return switch_obj(func)

maxframe/mixin.py

@@ -0,0 +1,100 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging.config
+from typing import Union
+
+import traitlets as T
+from tornado.log import enable_pretty_logging
+from traitlets.config import Configurable
+
+from . import env
+from .utils import trait_from_env
+
+
+class ServiceConfigMixin(Configurable):
+    port_env = env.MAXFRAME_SERVICE_PORT
+    port = T.Integer(
+        7953,
+        config=True,
+        help=f"Port on which to listen ({port_env} env var)",
+    )
+    port_default = trait_from_env("port", port_env)
+
+    port_retries_env = env.MAXFRAME_SERVICE_PORT_RETRIES
+    port_retries = T.Integer(
+        50,
+        config=True,
+        help=f"""Number of ports to try if the specified port is not available
+                           ({port_retries_env} env var)""",
+    )
+    port_retries_default = trait_from_env("port_retries", port_retries_env)
+
+    ip_env = env.MAXFRAME_SERVICE_LISTEN_ADDRESS
+    ip = T.Unicode(
+        "127.0.0.1",
+        config=True,
+        help=f"IP address on which to listen ({ip_env} env var)",
+    )
+    ip_default = trait_from_env("ip", ip_env)
+
+    allow_origin_env = env.MAXFRAME_SERVICE_ALLOW_ORIGIN
+    allow_origin = T.Unicode(
+        "",
+        config=True,
+        help=f"Sets the Access-Control-Allow-Origin header. ({allow_origin_env} env var)",
+    )
+    allow_origin_default = trait_from_env("allow_origin", allow_origin_env)
+
+    # Base URL
+    base_url_env = env.MAXFRAME_SERVICE_BASE_URL
+    base_url = T.Unicode(
+        "/",
+        config=True,
+        help=f"The misc path for mounting all API resources ({base_url_env} env var)",
+    )
+    base_url_default = trait_from_env("base_url", base_url_env)
+
+
+class LoggerConfigMixin(Configurable):
+    log_config_file_env = env.MAXFRAME_SERVICE_LOG_CONFIG_FILE
+    log_config_file = T.Unicode(
+        "",
+        config=True,
+        help=f"Sets the config file of logger. ({log_config_file_env} env var)",
+    )
+    log_config_file_default = trait_from_env("log_config_file", log_config_file_env)
+
+    def init_logger(self, log_level: Union[int, str] = None):
+        """
+        Init the logger. If the environment variable MAXFRAME_SERVICE_LOG_CONFIG_FILE is
+        set, the logger will be configured by this config file.
+        If log_level is set, it will override the level defined in the config file.
+
+        Parameters
+        ----------
+        log_level : int or str
+            The log level.
+        """
+        # Enable the same pretty logging the server uses
+        enable_pretty_logging()
+        if self.log_config_file:
+            logging.config.fileConfig(
+                self.log_config_file, disable_existing_loggers=False
+            )
+        # LogLevel from command line or environment has a high priority than config file.
+        if log_level:
+            logging.getLogger().setLevel(log_level)
+        # Adjust kubernetes logging level to hide secrets when logging
+        logging.getLogger("kubernetes").setLevel(logging.WARNING)

maxframe/odpsio/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .arrow import arrow_to_pandas, pandas_to_arrow
+from .schema import (
+    build_dataframe_table_meta,
+    odps_schema_to_pandas_dtypes,
+    pandas_to_odps_schema,
+)
+from .tableio import HaloTableIO

maxframe/odpsio/arrow.py

@@ -0,0 +1,91 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Tuple, Union
+
+import pandas as pd
+import pyarrow as pa
+
+from ..core import OutputType
+from ..protocol import DataFrameTableMeta
+from ..tensor.core import TENSOR_TYPE
+from ..typing_ import ArrowTableType, PandasObjectTypes
+from .schema import build_dataframe_table_meta
+
+
+def _rebuild_dataframe(
+    df: pd.DataFrame, table_meta: DataFrameTableMeta
+) -> Union[pd.DataFrame, pd.Series]:
+    indexed = df.set_index(list(table_meta.table_index_column_names))
+    indexed.index.set_names(table_meta.pd_index_level_names, inplace=True)
+    if len(table_meta.pd_column_level_names) == 1:
+        cols = pd.Index(
+            table_meta.pd_column_names, name=table_meta.pd_column_level_names[0]
+        )
+    else:
+        cols = pd.MultiIndex.from_tuples(table_meta.pd_column_names).set_names(
+            table_meta.pd_column_level_names
+        )
+    indexed.columns = cols
+    if table_meta.type == OutputType.series:
+        return indexed.iloc[:, 0]
+    return indexed
+
+
+def _rebuild_index(df: pd.DataFrame, table_meta: DataFrameTableMeta) -> pd.Index:
+    if df.shape[1] > 1:
+        df.columns = pd.Index(table_meta.pd_index_level_names)
+        return pd.MultiIndex.from_frame(df)
+    return pd.Index(df.iloc[:, 0], name=table_meta.pd_index_level_names[0])
+
+
+def arrow_to_pandas(
+    arrow_table: ArrowTableType, table_meta: DataFrameTableMeta
+) -> PandasObjectTypes:
+    df = arrow_table.to_pandas()
+    if table_meta.type in (OutputType.dataframe, OutputType.series):
+        return _rebuild_dataframe(df, table_meta)
+    elif table_meta.type == OutputType.index:
+        return _rebuild_index(df, table_meta)
+    elif table_meta.type == OutputType.scalar:
+        return _rebuild_index(df, table_meta)[0]
+    else:  # this could never happen  # pragma: no cover
+        raise ValueError(f"Does not support meta type {table_meta.type!r}")
+
+
+def pandas_to_arrow(
+    df: Any, nthreads=1, ignore_index=False
+) -> Tuple[ArrowTableType, DataFrameTableMeta]:
+    table_meta = build_dataframe_table_meta(df, ignore_index)
+    df = df.copy() if callable(getattr(df, "copy", None)) else df
+    if table_meta.type in (OutputType.dataframe, OutputType.series):
+        if table_meta.type == OutputType.series:
+            df = df.to_frame("_data" if df.name is None else df.name)
+        df.columns = pd.Index(table_meta.table_column_names)
+        if not ignore_index:
+            df = df.rename_axis(table_meta.table_index_column_names).reset_index()
+    elif ignore_index:
+        df = pd.DataFrame([], columns=[])
+    elif table_meta.type == OutputType.index:
+        names = [f"_idx_{idx}" for idx in range(len(df.names))]
+        df = df.to_frame(name=names[0] if len(names) == 1 else names)
+    elif table_meta.type == OutputType.scalar:
+        names = ["_idx_0"]
+        if isinstance(df, TENSOR_TYPE):
+            df = pd.DataFrame([], columns=names).astype({names[0]: df.dtype})
+        else:
+            df = pd.DataFrame([[df]], columns=names)
+    else:  # this could never happen  # pragma: no cover
+        raise ValueError(f"Does not support meta type {table_meta.type!r}")
+    return pa.Table.from_pandas(df, nthreads=nthreads, preserve_index=False), table_meta