wehub-resource-sync
156 lines
5.1 KiB
Python
156 lines
5.1 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import operator
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from .utils import hashable
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
BinaryOp = Callable[[Any, Any], Any]
|
|
UnaryOp = Callable[[Any], Any]
|
|
|
|
|
|
INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[BinaryOp, tuple[str, BinaryOp]] = {
|
|
# inplace op fn: (magic name, non-inplace op fn)
|
|
operator.iadd: ("__iadd__", operator.add),
|
|
operator.iand: ("__iand__", operator.and_),
|
|
operator.iconcat: ("__iconcat__", operator.concat),
|
|
operator.ifloordiv: ("__ifloordiv__", operator.floordiv),
|
|
operator.ilshift: ("__ilshift__", operator.lshift),
|
|
operator.imatmul: ("__imatmul__", operator.matmul),
|
|
operator.imod: ("__imod__", operator.mod),
|
|
operator.imul: ("__imul__", operator.mul),
|
|
operator.ior: ("__ior__", operator.or_),
|
|
operator.ipow: ("__ipow__", operator.pow),
|
|
operator.irshift: ("__irshift__", operator.rshift),
|
|
operator.isub: ("__isub__", operator.sub),
|
|
operator.itruediv: ("__itruediv__", operator.truediv),
|
|
operator.ixor: ("__ixor__", operator.xor),
|
|
}
|
|
|
|
NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[
|
|
BinaryOp, tuple[str, str | None]
|
|
] = {
|
|
# op fn: (magic name, reverse magic name)
|
|
operator.add: ("__add__", "__radd__"),
|
|
operator.and_: ("__and__", "__rand__"),
|
|
operator.contains: ("__contains__", None),
|
|
operator.delitem: ("__delitem__", None),
|
|
operator.eq: ("__eq__", "__eq__"),
|
|
operator.floordiv: ("__floordiv__", "__rfloordiv__"),
|
|
operator.ge: ("__ge__", "__le__"),
|
|
operator.getitem: ("__getitem__", None),
|
|
operator.gt: ("__gt__", "__lt__"),
|
|
operator.le: ("__le__", "__ge__"),
|
|
operator.lshift: ("__lshift__", "__rlshift__"),
|
|
operator.lt: ("__lt__", "__gt__"),
|
|
operator.matmul: ("__matmul__", "__rmatmul__"),
|
|
operator.mod: ("__mod__", "__rmod__"),
|
|
operator.mul: ("__mul__", "__rmul__"),
|
|
operator.ne: ("__ne__", "__ne__"),
|
|
operator.or_: ("__or__", "__ror__"),
|
|
operator.pow: ("__pow__", "__rpow__"),
|
|
operator.rshift: ("__rshift__", "__rrshift__"),
|
|
operator.sub: ("__sub__", "__rsub__"),
|
|
operator.truediv: ("__truediv__", "__rtruediv__"),
|
|
operator.xor: ("__xor__", "__rxor__"),
|
|
}
|
|
|
|
UNARY_OPS_TO_MAGIC_NAMES: dict[UnaryOp, str] = {
|
|
operator.neg: "__neg__",
|
|
operator.invert: "__invert__",
|
|
operator.pos: "__pos__",
|
|
operator.abs: "__abs__",
|
|
operator.index: "__index__",
|
|
operator.inv: "__inv__",
|
|
operator.truth: "__bool__",
|
|
bool: "__bool__",
|
|
abs: "__abs__",
|
|
float: "__float__",
|
|
len: "__len__",
|
|
int: "__int__",
|
|
complex: "__complex__",
|
|
}
|
|
# TODO(SigureMo): support any, all, sum
|
|
|
|
|
|
INPLACE_BINARY_OPS = set(INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
|
|
NON_INPLACE_BINARY_OPS = set(NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
|
|
BINARY_OPS = INPLACE_BINARY_OPS | NON_INPLACE_BINARY_OPS
|
|
UNARY_OPS = set(UNARY_OPS_TO_MAGIC_NAMES.keys())
|
|
|
|
|
|
# NOTE: Both operator.pow and operator.ipow should be considered for inclusion in this list,
|
|
# as they raise ZeroDivisionError when evaluating 0^n where n < 0 (division by zero).
|
|
NEED_GUARD_ZERO_DIVISION_ERROR_OPS: list[BinaryOp] = [
|
|
operator.floordiv,
|
|
operator.truediv,
|
|
operator.mod,
|
|
operator.ifloordiv,
|
|
operator.itruediv,
|
|
operator.imod,
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class MagicMethod:
|
|
name: str
|
|
is_inplace: bool = False
|
|
is_reverse: bool = False
|
|
|
|
|
|
def magic_method_builtin_dispatch(fn: BinaryOp | UnaryOp) -> list[MagicMethod]:
|
|
if not hashable(fn):
|
|
return []
|
|
if fn in INPLACE_BINARY_OPS:
|
|
inplace_magic_name, non_inplace_op = INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
|
|
fn
|
|
]
|
|
return [
|
|
MagicMethod(inplace_magic_name, is_inplace=True),
|
|
*magic_method_builtin_dispatch(non_inplace_op),
|
|
]
|
|
elif fn in NON_INPLACE_BINARY_OPS:
|
|
magic_name, reverse_magic_name = NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
|
|
fn
|
|
]
|
|
magic_methods = [MagicMethod(magic_name)]
|
|
if reverse_magic_name is not None:
|
|
magic_methods.append(
|
|
MagicMethod(reverse_magic_name, is_reverse=True)
|
|
)
|
|
return magic_methods
|
|
elif fn in UNARY_OPS:
|
|
magic_name = UNARY_OPS_TO_MAGIC_NAMES[fn]
|
|
return [MagicMethod(magic_name)]
|
|
return []
|
|
|
|
|
|
def non_inplace_op_to_inplace_op(
|
|
fn: BinaryOp,
|
|
) -> BinaryOp | None:
|
|
for inplace_op, (
|
|
_,
|
|
non_inplace_op,
|
|
) in INPLACE_BINARY_OPS_TO_MAGIC_NAMES.items():
|
|
if fn is non_inplace_op:
|
|
return inplace_op
|
|
return None
|