More features, refactoring

betterproto/templates/main.py

@@ -1,11 +1,30 @@
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
 # source: {{ description.filename }}
+# plugin: python-betterproto
+{% if description.enums %}import enum
+{% endif %}
 from dataclasses import dataclass
 from typing import List
 
 import betterproto
 
 
+{% if description.enums %}{% for enum in description.enums %}
+class {{ enum.name }}(enum.IntEnum):
+    {% if enum.comment %}
+{{ enum.comment }}
+
+    {% endif %}
+    {% for entry in enum.entries %}
+        {% if entry.comment %}
+{{ entry.comment }}
+        {% endif %}
+    {{ entry.name }} = {{ entry.value }}
+    {% endfor %}
+{% endfor %}
+
+
+{% endif %}
 {% for message in description.messages %}
 @dataclass
 class {{ message.name }}(betterproto.Message):