Merge pull request #41 from waszil/intenum

IntEum for python

Author: hit9

compiler/bitproto/renderer/impls/py/formatter.py

@@ -132,7 +132,7 @@ def format_default_value_message(self, t: Message) -> str:
         return f"{factory}()"
 
     def format_default_value_enum(self, t: Enum) -> str:
-        return "0"
+        return f"{self.format_type(t)}.{self.format_enum_field_name(t.fields()[0])}"
 
     def format_default_value_array(self, t: Array) -> str:
         cap = self.format_int_value(t.cap)

compiler/bitproto/renderer/impls/py/renderer.py

@@ -35,6 +35,8 @@
 from bitproto.renderer.renderer import Renderer
 from bitproto.utils import cached_property, override
 
+_enum_field_proxy_prefix = "_enum_field_proxy__"
+
 
 class BlockProtoDocstring(BlockBindProto[F]):
     @override(Block)
@@ -47,7 +49,8 @@ class BlockGeneralImports(Block[F]):
     def render(self) -> None:
         self.push("import json")
         self.push("from dataclasses import dataclass, field")
-        self.push("from typing import ClassVar, Dict, List")
+        self.push("from typing import ClassVar, Dict, List, Union")
+        self.push("from enum import IntEnum, unique")
         self.push_empty_line()
         self.push("from bitprotolib import bp")
 
@@ -140,7 +143,7 @@ class BlockEnumField(BlockBindEnumField[F]):
     def render(self) -> None:
         self.push_definition_comments()
         self.push(
-            f"{self.enum_field_name}: {self.enum_field_type} = {self.enum_field_value}"
+            f"{self.enum_field_name}: {self.enum_field_type} = {self.enum_field_type}.{self.enum_field_name}"
         )
 
 
@@ -159,20 +162,58 @@ class BlockEnumFieldListWrapper(BlockBindEnum[F], BlockWrapper[F]):
     def wraps(self) -> Block:
         return BlockEnumFieldList(self.d)
 
+    @override(BlockWrapper)
+    def before(self) -> None:
+        self.push_comment("Aliases for backwards compatibility")
+
+    @override(BlockWrapper)
+    def after(self) -> None:
+        self.push_empty_line()
+
+
+class BlockIntEnumField(BlockBindEnumField[F]):
+    @override(Block)
+    def render(self) -> None:
+        self.push_definition_comments()
+        self.push(f"    {self.enum_field_name} = {self.enum_field_value}")
+
+
+class BlockIntEnumFieldList(BlockBindEnum[F], BlockComposition[F]):
+    @override(BlockComposition)
+    def blocks(self) -> List[Block]:
+        return [BlockIntEnumField(field) for field in self.d.fields()]
+
+    @override(BlockComposition)
+    def separator(self) -> str:
+        return "\n"
+
+
+class BlockIntEnumFieldListWrapper(BlockBindEnum[F], BlockWrapper[F]):
+    @override(BlockWrapper)
+    def wraps(self) -> Block:
+        return BlockIntEnumFieldList(self.d)
+
     @override(BlockWrapper)
     def before(self) -> None:
         self.render_enum_type()
 
+    @override(BlockWrapper)
+    def after(self) -> None:
+        self.push_empty_line()
+
     def render_enum_type(self) -> None:
         self.push_definition_comments()
-        self.push(f"{self.enum_name} = int")
+        self.push("@unique")
+        self.push(f"class {self.enum_name}(IntEnum):")
         self.push_typing_hint_inline_comment()
 
 
 class BlockEnumValueToNameMapItem(BlockBindEnumField[F]):
     @override(Block)
     def render(self) -> None:
-        self.push(f'{self.enum_field_value}: "{self.enum_field_name}",')
+        self.push(
+            f'{self.enum_field_type}.{self.enum_field_name}: "{self.enum_field_name}",'
+        )
 
 
 class BlockEnumValueToNameMapItemList(BlockBindEnum[F], BlockComposition[F]):
@@ -215,6 +256,7 @@ class BlockEnum(BlockBindEnum[F], BlockComposition[F]):
     @override(BlockComposition[F])
     def blocks(self) -> List[Block]:
         return [
+            BlockIntEnumFieldListWrapper(self.d),
             BlockEnumFieldListWrapper(self.d),
             BlockEnumValueToNameMap(self.d),
             BlockEnumMethodProcessor(self.d),
@@ -229,9 +271,22 @@ def message_field_default_value(self) -> str:
     @override(Block)
     def render(self) -> None:
         self.push_definition_comments()
-        self.push(
-            f"{self.message_field_name}: {self.message_field_type} = {self.message_field_default_value}"
-        )
+        if issubclass(type(self.d.type), Enum):
+            # use union of int and IntEnum so that field assignment works as before without typing problem.
+            self.push(
+                f"{self.message_field_name}: Union[int, {self.message_field_type}] = {self.message_field_default_value}"
+            )
+            # push a proxy field for enum  fields to hold integer value
+            self.push_comment(
+                f"This field is a proxy to hold integer value of enum field '{self.message_field_name}'"
+            )
+            self.push(
+                f"{_enum_field_proxy_prefix}{self.message_field_name}: int = field(init=False, repr=False)"
+            )
+        else:
+            self.push(
+                f"{self.message_field_name}: {self.message_field_type} = {self.message_field_default_value}"
+            )
         self.push_typing_hint_inline_comment()
 
 
@@ -289,6 +344,19 @@ def before(self) -> None:
         self.push_definition_docstring(indent=4)
 
 
+class BlockMessageDictFactory(BlockMessageBase, BlockWrapper[F]):
+    def wraps(self) -> Block[F]:
+        pass
+
+    @override(BlockWrapper)
+    def before(self) -> None:
+        self.push("@staticmethod")
+        self.push("def dict_factory(kv_pairs):")
+        self.push(
+            f"    return {{k: v for k, v in kv_pairs if not k.startswith('{_enum_field_proxy_prefix}')}}"
+        )
+
+
 class BlockMessageMethodProcessorFieldItem(BlockBindMessageField[F]):
     @override(Block)
     def render(self) -> None:
@@ -310,6 +378,104 @@ def separator(self) -> str:
         return "\n"
 
 
+class BlockMessagePostInitField(BlockBindMessageField[F]):
+    @override(Block)
+    def render(self) -> None:
+        self.push_comment(
+            f"initialize handling of enum field '{self.message_field_name}' as `enum.IntEnum`"
+        )
+        self.push(
+            f'if not isinstance(getattr({self.d.message.name}, "{self.message_field_name}", False), property):'
+        )
+        self.push(
+            f"    self.{_enum_field_proxy_prefix}{self.message_field_name} = self.{self.message_field_name}"
+        )
+        self.push(
+            f"    {self.d.message.name}.{self.message_field_name} = property("
+            f"{self.d.message.name}._get_{self.message_field_name}, "
+            f"{self.d.message.name}._set_{self.message_field_name})"
+        )
+
+
+class BlockMessagePostInitPass(BlockBindMessage[F]):
+    @override(Block)
+    def render(self) -> None:
+        self.push("pass")
+
+
+class BlockMessagePostInitItemList(BlockMessageBase, BlockComposition[F]):
+    @override(BlockComposition)
+    def blocks(self) -> List[Block[F]]:
+        if any(
+            (issubclass(type(field.type), Enum) for field in self.d.sorted_fields())
+        ):
+            b: List[Block[F]] = [
+                BlockMessagePostInitField(field, indent=self.indent)
+                for field in self.d.sorted_fields()
+                if issubclass(type(field.type), Enum)
+            ]
+            return b
+        else:
+            return [BlockMessagePostInitPass(self.d, indent=self.indent)]
+
+    @override(BlockComposition)
+    def separator(self) -> str:
+        return "\n"
+
+
+class BlockMessagePostInit(BlockBindMessage[F], BlockWrapper[F]):
+    @override(BlockWrapper)
+    def wraps(self) -> Block[F]:
+        return BlockMessagePostInitItemList(self.d, indent=self.indent + 4)
+
+    @override(BlockWrapper)
+    def before(self) -> None:
+        self.push("def __post_init__(self):")
+
+
+class BlockMessageEnumProxyFieldAccessorField(BlockBindMessageField[F]):
+    @override(Block)
+    def render(self) -> None:
+        self.push(
+            f"def _get_{self.message_field_name}(self) -> {self.message_field_type}:"
+        )
+        self.push(f'    """property getter for enum proxy field"""')
+        self.push(
+            f"    return {self.message_field_type}(self.{_enum_field_proxy_prefix}{self.message_field_name})\n"
+        )
+        self.push(f"def _set_{self.message_field_name}(self, val):")
+        self.push(f'    """property setter for enum proxy field"""')
+        self.push(f"    self.{_enum_field_proxy_prefix}{self.message_field_name} = val")
+
+
+class BlockMessageEnumProxyFieldAccessorFieldList(
+    BlockMessageBase, BlockComposition[F]
+):
+    @override(BlockComposition)
+    def blocks(self) -> List[Block[F]]:
+        if any(
+            (issubclass(type(field.type), Enum) for field in self.d.sorted_fields())
+        ):
+            b: List[Block[F]] = [
+                BlockMessageEnumProxyFieldAccessorField(field, indent=self.indent)
+                for field in self.d.sorted_fields()
+                if issubclass(type(field.type), Enum)
+            ]
+            return b
+        else:
+            return []
+
+    @override(BlockComposition)
+    def separator(self) -> str:
+        return "\n"
+
+
+class BlockMessageEnumProxyFieldAccessors(BlockBindMessage[F], BlockWrapper[F]):
+    @override(BlockWrapper)
+    def wraps(self) -> Block[F]:
+        return BlockMessageEnumProxyFieldAccessorFieldList(self.d, indent=self.indent)
+
+
 class BlockMessageMethodProcessor(BlockMessageBase, BlockWrapper[F]):
     @override(BlockWrapper)
     def wraps(self) -> Block[F]:
@@ -606,6 +772,9 @@ class BlockMessage(BlockMessageBase, BlockComposition[F]):
     def blocks(self) -> List[Block[F]]:
         return [
             BlockMessageClass(self.d),
+            BlockMessagePostInit(self.d, indent=4),
+            BlockMessageDictFactory(self.d, indent=4),
+            BlockMessageEnumProxyFieldAccessors(self.d, indent=4),
             BlockMessageMethodProcessor(self.d, indent=4),
             BlockMessageMethodSetByte(self.d, indent=4),
             BlockMessageMethodGetByte(self.d, indent=4),

docs/language.rst

@@ -651,3 +651,9 @@ Vim
 A syntax plugin for `vim <https://www.vim.org/>`_ is available from
 `bitproto's github repository <https://github.com/hit9/bitproto/tree/master/editors/vim>`_.
 This plugin only supports syntax highlighting of bitproto language.
+
+PyCharm
+"""""""
+
+Syntax highlighting settings for PyCharm is available from
+`bitproto's github repository <https://github.com/hit9/bitproto/tree/master/editors/pycharm>`_.

docs/python-guide.rst

@@ -30,16 +30,22 @@ the mapped classes, constants and api methods etc.
 
 In the generated ``pen_bp.py``:
 
-* The ``enum Color`` in bitproto is mapped to a typing hint alias, and the enum values are mapped
-  to constants:
+* The ``enum Color`` in bitproto is mapped to an ``enum.IntEum`` class, and the enum values are mapped
+  to constants as well:
 
   .. sourcecode:: python
 
-     Color = int # 3bit
-     COLOR_UNKNOWN: Color = 0
-     COLOR_RED: Color = 1
-     COLOR_BLUE: Color = 2
-     COLOR_GREEN: Color = 3
+     @unique
+     class Color(IntEnum): # 3 bit
+        COLOR_UNKNOWN = 0
+        COLOR_RED = 1
+        COLOR_BLUE = 2
+        COLOR_GREEN = 3
+
+     COLOR_UNKNOWN: Color = Color.COLOR_UNKNOWN
+     COLOR_RED: Color = Color.COLOR_RED
+     COLOR_BLUE: Color = Color.COLOR_BLUE
+     COLOR_GREEN: Color = Color.COLOR_GREEN
 
 * The ``Timestamp`` in bitproto is also mapped to a typing hint alias:
 

editors/pycharm/README.rst

@@ -0,0 +1,11 @@
+PyCharm syntax highlighting for bitproto
+========================================
+
+PyCharm syntax highlighting option for bitproto.
+
+Installation
+------------
+
+File -> Manage IDE Settings -> Import Settings
+
+and browse settings.zip in this folder.

editors/pycharm/settings.zip

@@ -181,7 +181,7 @@ class MessageBase(Accessor):
 
     def to_dict(self) -> Dict[str, Any]:
         """Converts this message to a dict."""
-        return asdict(self)
+        return asdict(self, dict_factory=getattr(self, "dict_factory", dict))
 
     def to_json(
         self, indent: Optional[int] = None, separators: Optional[Tuple[str, str]] = None

lib/py/bitprotolib/bp.py

@@ -0,0 +1,52 @@
+NAME=enums
+BIN=main
+
+BP_FILENAME=$(NAME).bitproto
+BP_C_FILENAME=$(NAME)_bp.c
+BP_GO_FILENAME=$(NAME)_bp.go
+BP_PY_FILENAME=$(NAME)_bp.py
+BP_LIB_DIR=../../../../../lib/c
+BP_LIC_C_PATH=$(BP_LIB_DIR)/bitproto.c
+
+C_SOURCE_FILE=main.c
+C_SOURCE_FILE_LIST=$(C_SOURCE_FILE) $(BP_C_FILENAME) $(BP_LIC_C_PATH) 
+C_BIN=$(BIN)
+
+GO_BIN=$(BIN)
+
+PY_SOURCE_FILE=main.py
+
+CC_OPTIMIZATION_ARG?=
+
+OPTIMIZATION_MODE_ARGS?=
+
+bp-c:
+	@bitproto c $(BP_FILENAME) c/ $(OPTIMIZATION_MODE_ARGS)
+
+bp-go:
+	@bitproto go $(BP_FILENAME) go/bp/ $(OPTIMIZATION_MODE_ARGS)
+
+bp-py:
+	@bitproto py $(BP_FILENAME) py/
+
+build-c: bp-c
+	@cd c && $(CC) $(C_SOURCE_FILE_LIST) -I. -I$(BP_LIB_DIR) -o $(C_BIN) $(CC_OPTIMIZATION_ARG)
+
+build-go: bp-go
+	@cd go && go build -o $(GO_BIN)
+
+build-py: bp-py
+
+run-c: build-c
+	@cd c && ./$(C_BIN)
+
+run-go: build-go
+	@cd go && ./$(GO_BIN)
+
+run-py: build-py
+	@cd py && python $(PY_SOURCE_FILE)
+
+clean:
+	@rm -fr c/$(C_BIN) go/$(GO_BIN) go/vendor */*_bp.* */**/*_bp.* py/__pycache__
+
+run: run-c run-go run-py

tests/test_encoding/encoding-cases/enums/Makefile

@@ -0,0 +1,4 @@
+---
+BasedOnStyle: Google
+IndentWidth: 4
+---