Gecko [ betsfootball.net ]

Name	Size	Permission	Date
__pycache__	[ DIR ]	drwxr-xr-x	2025-04-08 17:01
integration	[ DIR ]	drwxr-xr-x	2025-04-08 17:01
_C.cpython-311-x86_64-linux-gnu.so	796.47 KB	-rwxr-xr-x	2025-04-08 17:01
_C.pyi	6 KB	-rw-r--r--	2025-04-08 17:01
__init__.py	5.75 KB	-rw-r--r--	2025-04-08 17:01
accessor.py	14.5 KB	-rw-r--r--	2025-04-08 17:01
dataclasses.py	17.98 KB	-rw-r--r--	2025-04-08 17:01
functools.py	6.26 KB	-rw-r--r--	2025-04-08 17:01
ops.py	149.48 KB	-rw-r--r--	2025-04-08 17:01
py.typed	0 B	-rw-r--r--	2025-04-08 17:01
pytree.py	4.24 KB	-rw-r--r--	2025-04-08 17:01
registry.py	29.54 KB	-rw-r--r--	2025-04-08 17:01
treespec.py	1.88 KB	-rw-r--r--	2025-04-08 17:01
typing.py	18.56 KB	-rw-r--r--	2025-04-08 17:01
utils.py	3.42 KB	-rw-r--r--	2025-04-08 17:01
version.py	1.77 KB	-rw-r--r--	2025-04-08 17:01

Rename

# Copyright 2022-2025 MetaOPT Team. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Utility functions for OpTree."""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Callable, Iterable, overload

if TYPE_CHECKING:
    from optree.typing import S, T, U

def total_order_sorted(
    iterable: Iterable[T],
    /,
    *,
    key: Callable[[T], Any] | None = None,
    reverse: bool = False,
) -> list[T]:
    """Sort an iterable in a total order.

This is useful for sorting objects that are not comparable, e.g., dictionaries with different
    types of keys.
    """
    sequence = list(iterable)

try:
        # Sort directly if possible
        return sorted(sequence, key=key, reverse=reverse)  # type: ignore[type-var,arg-type]
    except TypeError:
        if key is None:

def key_fn(x: T) -> tuple[str, Any]:
                return (f'{x.__class__.__module__}.{x.__class__.__qualname__}', x)

else:

def key_fn(x: T) -> tuple[str, Any]:
                y = key(x)
                return (f'{y.__class__.__module__}.{y.__class__.__qualname__}', y)

try:
            # Add `{obj.__class__.__module__}.{obj.__class__.__qualname__}` to the key order to make
            # it sortable between different types (e.g., `int` vs. `str`)
            return sorted(sequence, key=key_fn, reverse=reverse)
        except TypeError:  # cannot sort the keys (e.g., user-defined types)
            return sequence  # fallback to original order

@overload
def safe_zip(
    iter1: Iterable[T],
    /,
) -> zip[tuple[T]]: ...

@overload
def safe_zip(
    iter1: Iterable[T],
    iter2: Iterable[S],
    /,
) -> zip[tuple[T, S]]: ...

@overload
def safe_zip(
    iter1: Iterable[T],
    iter2: Iterable[S],
    iter3: Iterable[U],
    /,
) -> zip[tuple[T, S, U]]: ...

@overload
def safe_zip(
    iter1: Iterable[Any],
    iter2: Iterable[Any],
    iter3: Iterable[Any],
    iter4: Iterable[Any],
    /,
    *iters: Iterable[Any],
) -> zip[tuple[Any, ...]]: ...

def safe_zip(*args: Iterable[Any]) -> zip[tuple[Any, ...]]:
    """Strict zip that requires all arguments to be the same length."""
    seqs = [arg if isinstance(arg, Sequence) else list(arg) for arg in args]
    if len(set(map(len, seqs))) > 1:
        raise ValueError(f'length mismatch: {list(map(len, seqs))}')
    return zip(*seqs)

def unzip2(xys: Iterable[tuple[T, S]], /) -> tuple[tuple[T, ...], tuple[S, ...]]:
    """Unzip sequence of length-2 tuples into two tuples."""
    # Note: we deliberately don't use zip(*xys) because it is lazily evaluated,
    # is too permissive about inputs, and does not guarantee a length-2 output.
    # For example, for empty dict: tuple(zip(*{}.items())) -> ()
    xs = []
    ys = []
    for x, y in xys:
        xs.append(x)
        ys.append(y)
    return tuple(xs), tuple(ys)