149 lines
No EOL
4.8 KiB
Python
149 lines
No EOL
4.8 KiB
Python
"""
|
|
Two magic tricks for classes:
|
|
|
|
class X(metaclass=extendabletype):
|
|
...
|
|
|
|
# in some other file...
|
|
class __extend__(X):
|
|
... # add new methods and class attributes to X
|
|
|
|
Mostly useful together with the second trick, which lets you build
|
|
methods whose 'self' is a pair of objects instead of just one:
|
|
|
|
class __extend__(pairtype(X, Y)):
|
|
attribute = 42
|
|
def method(self, other, arguments):
|
|
x, y = self
|
|
...
|
|
|
|
pair(x, y).attribute
|
|
pair(x, y).method(other, arguments)
|
|
|
|
This finds methods and class attributes based on the actual
|
|
class of both objects that go into the pair(), with the usual
|
|
rules of method/attribute overriding in (pairs of) subclasses.
|
|
|
|
For more information, see test_pairtype.
|
|
"""
|
|
|
|
from typing import Any, Dict, Tuple, Type, TypeVar, Generic, Callable, Optional, Iterable
|
|
|
|
T1 = TypeVar('T1')
|
|
T2 = TypeVar('T2')
|
|
|
|
class extendabletype(type):
|
|
"""A type with a syntax trick: 'class __extend__(t)' actually extends
|
|
the definition of 't' instead of creating a new subclass."""
|
|
def __new__(mcs, name: str, bases: Tuple[type, ...], dct: Dict[str, Any]) -> Optional[type]:
|
|
if name == '__extend__':
|
|
for cls in bases:
|
|
for key, value in dct.items():
|
|
if key == '__module__':
|
|
continue
|
|
# Add attributes to the base class
|
|
setattr(cls, key, value)
|
|
return None
|
|
else:
|
|
return super().__new__(mcs, name, bases, dct)
|
|
|
|
def pair(a: T1, b: T2) -> Tuple[T1, T2]:
|
|
"""Return a pair object with dynamic type dispatch."""
|
|
tp = pairtype(type(a), type(b))
|
|
return tp((a, b)) # tp is a subclass of tuple
|
|
|
|
pairtypecache: Dict[Tuple[type, type], type] = {}
|
|
|
|
def pairtype(cls1: type, cls2: type) -> type:
|
|
"""type(pair(a,b)) is pairtype(a.__class__, b.__class__)."""
|
|
if (cls1, cls2) in pairtypecache:
|
|
return pairtypecache[(cls1, cls2)]
|
|
|
|
# Generate a meaningful name for the new pair type
|
|
name = f'pairtype({cls1.__name__}, {cls2.__name__})'
|
|
|
|
# Create base types for the new pair type
|
|
bases1 = [pairtype(base1, cls2) for base1 in cls1.__bases__]
|
|
bases2 = [pairtype(cls1, base2) for base2 in cls2.__bases__]
|
|
bases = tuple(bases1 + bases2) or (tuple,) # 'tuple': ultimate base
|
|
|
|
# Create the new pair type
|
|
pair = pairtypecache[(cls1, cls2)] = extendabletype(name, bases, {})
|
|
return pair
|
|
|
|
def pairmro(cls1: type, cls2: type) -> Iterable[Tuple[type, type]]:
|
|
"""
|
|
Return the resolution order on pairs of types for double dispatch.
|
|
|
|
This order is compatible with the mro of pairtype(cls1, cls2).
|
|
"""
|
|
for base2 in cls2.__mro__:
|
|
for base1 in cls1.__mro__:
|
|
yield base1, base2
|
|
|
|
class DoubleDispatchRegistry:
|
|
"""
|
|
A mapping of pairs of types to arbitrary objects respecting inheritance
|
|
"""
|
|
def __init__(self):
|
|
self._registry: Dict[Tuple[type, type], Callable] = {}
|
|
self._cache: Dict[Tuple[type, type], Callable] = {}
|
|
|
|
def __getitem__(self, clspair: Tuple[type, type]) -> Callable:
|
|
cls1, cls2 = clspair
|
|
# Check cache first
|
|
if clspair in self._cache:
|
|
return self._cache[clspair]
|
|
|
|
# Traverse the MRO to find the closest match
|
|
for c1, c2 in pairmro(cls1, cls2):
|
|
if (c1, c2) in self._cache:
|
|
return self._cache[(c1, c2)]
|
|
|
|
# If no match found, use the default implementation
|
|
return self._registry.get(clspair, None)
|
|
|
|
def __setitem__(self, clspair: Tuple[type, type], value: Callable):
|
|
self._registry[clspair] = value
|
|
self._cache = self._registry.copy()
|
|
|
|
class DoubleDispatchFunction:
|
|
def __init__(self, default_func: Callable):
|
|
self._registry = DoubleDispatchRegistry()
|
|
self._default = default_func
|
|
|
|
def __call__(self, arg1: Any, arg2: Any, *args, **kwargs) -> Any:
|
|
func = self._registry[(type(arg1), type(arg2))]
|
|
if func is None:
|
|
func = self._default
|
|
return func(arg1, arg2, *args, **kwargs)
|
|
|
|
def register(self, cls1: Type, cls2: Type) -> Callable:
|
|
def decorator(func: Callable) -> Callable:
|
|
self._registry[(cls1, cls2)] = func
|
|
return func
|
|
return decorator
|
|
|
|
def doubledispatch(default_func: Optional[Callable] = None) -> Callable:
|
|
"""
|
|
Decorator returning a double-dispatch function
|
|
|
|
Usage
|
|
-----
|
|
@doubledispatch
|
|
def func(x, y):
|
|
return 0
|
|
|
|
@func.register(str, str)
|
|
def func_string_string(x, y):
|
|
return 42
|
|
|
|
func(1, 2) # returns 0
|
|
func('x', 'y') # returns 42
|
|
"""
|
|
def decorator(func):
|
|
return DoubleDispatchFunction(func)
|
|
|
|
if default_func is None:
|
|
return decorator
|
|
return decorator(default_func) |