hit9 / bitproto

The bit level data interchange format for serializing data structures (long term maintenance).
https://bitproto.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
127 stars 16 forks source link

Representing enums as enum.IntEnum #40

Closed waszil closed 2 years ago

waszil commented 2 years ago

Hi, great library! I have one question: do you plan to implement enum type representation with enum.IntEnum? It would be nice to be able to use enum classes instead of simple constants.

Example proto:

proto enums;

enum MyEnum : uint2 {
    MY_ENUM_UNKNOWN = 0;
    MY_ENUM_ONE = 1;
    MY_ENUM_TWO = 2;
    MY_ENUM_THREE = 3;
}

message EnumContainer {
    MyEnum my_enum = 2;
}

If the generated python code would be something like this:

import json
from dataclasses import dataclass, field
from typing import ClassVar, Dict, List
from enum import IntEnum, unique

from bitprotolib import bp

@unique
class MyEnum(IntEnum):
    MY_ENUM_UNKNOWN = 0
    MY_ENUM_ONE = 1
    MY_ENUM_TWO = 2
    MY_ENUM_THREE = 3

_MYENUM_VALUE_TO_NAME_MAP: Dict[MyEnum, str] = {
    MyEnum.MY_ENUM_UNKNOWN: "MY_ENUM_UNKNOWN",
    MyEnum.MY_ENUM_ONE: "MY_ENUM_ONE",
    MyEnum.MY_ENUM_TWO: "MY_ENUM_TWO",
    MyEnum.MY_ENUM_THREE: "MY_ENUM_THREE",
}

def bp_processor_MyEnum() -> bp.Processor:
    return bp.EnumProcessor(bp.Uint(2))

@dataclass
class EnumContainer(bp.MessageBase):
    # Number of bytes to serialize class EnumContainer
    BYTES_LENGTH: ClassVar[int] = 1

    # based on https://stackoverflow.com/a/61709025/1169220
    my_enum: MyEnum = MyEnum.MY_ENUM_UNKNOWN
    _my_enum: int = field(init=False, repr=False)

    def __post_init__(self):
        if not isinstance(getattr(EnumContainer, "my_enum", False), property):
            self._my_enum = self.my_enum
            EnumContainer.my_enum = property(EnumContainer._get_my_enum, EnumContainer._set_my_enum)

    def _get_my_enum(self):
        return MyEnum(self._my_enum)

    def _set_my_enum(self, val):
        self._my_enum = val

...

Then it could be used like this:

import enums_bp as bp

enum_container = bp.EnumContainer(my_enum=bp.MyEnum.MY_ENUM_ONE)
s = enum_container.encode()
enum_container_new = bp.EnumContainer()
enum_container_new.decode(s)
assert enum_container_new.my_enum == enum_container.my_enum
assert enum_container_new.encode() == s
assert isinstance(enum_container.my_enum, bp.MyEnum)
assert isinstance(enum_container_new.my_enum, bp.MyEnum)
assert enum_container.my_enum is bp.MyEnum.MY_ENUM_ONE
assert enum_container_new.my_enum is bp.MyEnum.MY_ENUM_ONE

What do you think?

hit9 commented 2 years ago

Sounds good.

But I don't have much time to do this lately...

Anyway, pull requests are welcome if you are waiting it urgently.

waszil commented 2 years ago

I'm working on it :)

waszil commented 2 years ago

PR #41 created

waszil commented 2 years ago

Thanks for the merge!