C++ Scoped Enums implementation

README.md

C++ Scoped Enums Implementation

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.

scoped_enum.py

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