This is an implementation for my C++ Scoped Enums blogpost.
The names are a bit different than in the post, but the ideas are the same and should map across.
This is an implementation for my C++ Scoped Enums blogpost.
The names are a bit different than in the post, but the ideas are the same and should map across.
from __future__ import annotations | |
import collections | |
import inspect | |
from collections import Counter, ChainMap | |
import attrs | |
import operator | |
from typing import Any | |
@attrs.frozen | |
class Get: | |
name: str | |
@attrs.frozen | |
class Set: | |
name: str | |
value: Proxy | |
@attrs.frozen | |
class Proxy: | |
lhs: Any | |
op: Any | |
rhs: Any | |
names: Counter[str] = attrs.field(factory=lambda: Counter()) | |
@staticmethod | |
def new(lhs, op, rhs) -> Proxy: | |
names = Counter() | |
if isinstance(lhs, Proxy): | |
names.update(lhs.names) | |
if isinstance(rhs, Proxy): | |
names.update(rhs.names) | |
return Proxy(lhs=lhs, op=op, rhs=rhs, names=names) | |
@staticmethod | |
def placeholder(name)->Proxy: | |
return Proxy(name, None, None, Counter([name])) | |
def _apply_operand(self, operand, resolve_name): | |
if isinstance(operand, Proxy): | |
return operand.apply(resolve_name) | |
return operand | |
def apply(self, resolve_name): | |
if self.op is None: | |
return resolve_name(self.lhs) | |
lhs = self._apply_operand(self.lhs, resolve_name) | |
rhs = self._apply_operand(self.rhs, resolve_name) | |
return self.op(lhs, rhs) | |
def __add__(self, other): | |
return Proxy.new(self, operator.add, other) | |
def __radd__(self, other): | |
return Proxy.new(other, operator.add, self) | |
def __sub__(self, other): | |
return Proxy.new(self, operator.sub, other) | |
def __rsub__(self, other): | |
return Proxy.new(other, operator.sub, self) | |
def __mul__(self, other): | |
return Proxy.new(self, operator.mul, other) | |
def __rmul__(self, other): | |
return Proxy.new(other, operator.mul, self) | |
def __truediv__(self, other): | |
return Proxy.new(self, operator.truediv, other) | |
def __rtruediv__(self, other): | |
return Proxy.new(other, operator.truediv, self) | |
def __floordiv__(self, other): | |
return Proxy.new(self, operator.floordiv, other) | |
def __rfloordiv__(self, other): | |
return Proxy.new(other, operator.floordiv, self) | |
def __pow__(self, other, modulo=None): | |
return Proxy.new(self, operator.pow, other) | |
def __rpow__(self, other): | |
return Proxy.new(other, operator.pow, self) | |
def __lshift__(self, other): | |
return Proxy.new(self, operator.lshift, other) | |
def __rlshift__(self, other): | |
return Proxy.new(other, operator.lshift, self) | |
def __rshift__(self, other): | |
return Proxy.new(self, operator.rshift, other) | |
def __rrshift__(self, other): | |
return Proxy.new(other, operator.rshift, self) | |
class Namespace(dict): | |
def __init__(self): | |
super().__init__() | |
self.log = [] | |
def __setitem__(self, name, value: Proxy): | |
if name.startswith("__") and name.endswith("__"): | |
return super().__setitem__(name, value) | |
self.log.append(Set(name, value)) | |
def __getitem__(self, name): | |
if name.startswith("__") and name.endswith("__"): | |
return super().__getitem__(name) | |
self.log.append(Get(name)) | |
return Proxy.placeholder(name=name) | |
@attrs.frozen | |
class Define: | |
name: str | |
@attrs.frozen | |
class Assign: | |
target: str | |
source: Proxy | |
def commands_to_actions(get_cmds: list[Get], set_cmd: Set | None): | |
if not set_cmd: | |
for cmd in get_cmds: | |
yield Define(cmd.name) | |
return | |
if isinstance(set_cmd.value, Proxy): | |
assigned_from = set_cmd.value.names | |
else: | |
assigned_from = Counter() | |
get_count = collections.Counter(get.name for get in get_cmds) | |
for name, count in get_count.items(): | |
if count > assigned_from[name]: | |
yield Define(name) | |
yield Assign(set_cmd.name, set_cmd.value) | |
def construct_dict(actions, namespace): | |
dct = {} | |
last_value = -1 | |
def resolve_name(name: str) -> Any: | |
if name in dct: | |
return dct[name] | |
placeholder = object() | |
value = namespace.get(name, placeholder) | |
if value == placeholder: | |
raise NameError() | |
return value | |
for action in actions: | |
match action: | |
case Define(name=name): | |
last_value += 1 | |
dct[name] = last_value | |
case Assign(target=name, source=value): | |
if isinstance(value, Proxy): | |
last_value = value.apply(resolve_name) | |
else: | |
last_value = value | |
dct[name] = last_value | |
return dct | |
class ScopedEnumMeta(type): | |
@classmethod | |
def __prepare__(metacls, name, bases): | |
# Return our custom namespace object | |
return Namespace() | |
def __new__(cls, name, bases, classdict): | |
# Convert the custom object to a regular dict, to avoid unwanted shenanigans. | |
log = iter(classdict.log) | |
actions = [] | |
get_cmds = [] | |
for cmd in log: | |
if isinstance(cmd, Get): | |
get_cmds.append(cmd) | |
continue | |
actions.extend(commands_to_actions(get_cmds, cmd)) | |
get_cmds = [] | |
actions.extend(commands_to_actions(get_cmds, None)) | |
frame = inspect.stack()[1].frame | |
enum_dict = construct_dict( | |
actions, ChainMap(frame.f_locals, frame.f_globals, frame.f_builtins) | |
) | |
classdict = dict(classdict) | enum_dict | |
return type.__new__(cls, name, bases, classdict) | |
class ScopedEnum(metaclass=ScopedEnumMeta): | |
pass |