Skip to content

Commit

Permalink
feat: render enum as (str, enum.Enum) (#637) (#652)
Browse files Browse the repository at this point in the history 
* feat: render enum as (str, enum.Enum) (#637)

* docs: fix typo in docs

* fix(test): add missing newline in expected_result

* feat: support (str, enum.Enum) and (enum.Enum)

To use the @enum.nonmember decorator, we need Python >= 3.11.
These changes aim at supporting (str, enum.Enum) and the use of @enum.nonmember for >= 3.11.
But also to support >=3.8,<3.11 we omit @enum.nonmembmer and only use (enum.Enum).

* docs: support >=3.11 (str, enum.Enum) and >=3.8,<3.11 (enum.Enum)

---------

Co-authored-by: Marcos Schroh <[email protected]>
  • Loading branch information
@alxdrcirilo @marcosschroh
alxdrcirilo and marcosschroh authored Jul 10, 2024
1 parent dd109c1 commit 18e226e
Show file tree
Hide file tree
Showing 13 changed files with 780 additions and 351 deletions.
2 changes: 1 addition & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ import typing
from dataclasses_avroschema import AvroModel, types


class FavoriteColor(enum.Enum):
class FavoriteColor(str, enum.Enum):
BLUE = "BLUE"
YELLOW = "YELLOW"
GREEN = "GREEN"
Expand Down
10 changes: 7 additions & 3 deletions dataclasses_avroschema/fields/fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -455,15 +455,19 @@ def __post_init__(self) -> None:
def _get_meta_class_attributes(self) -> typing.Dict[str, typing.Any]:
# get Enum members
members = self.type.__members__
meta = members.get("Meta")
meta = members.get("Meta") or getattr(self.type, "Meta", None)

doc: typing.Optional[str] = self.type.__doc__

# On python < 3.11 Enums have a default documentation so we remove it
if version.PY_VERSION < (3, 11) and doc == "An enumeration.":
doc = None

if meta is not None:
meta = meta.value
try:
if meta is not None:
meta = meta.value
except AttributeError:
pass

metadata = utils.FieldMetadata.create(meta)
if doc is not None:
Expand Down
18 changes: 10 additions & 8 deletions dataclasses_avroschema/model_generator/lang/python/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def render_extras(self) -> str:
return "".join([extra for extra in self.extras])

def render_metaclass(
self, *, schema: JsonDict, field_order: typing.Optional[typing.List[str]] = None
self, *, schema: JsonDict, field_order: typing.Optional[typing.List[str]] = None, decorator: str = ""
) -> typing.Optional[str]:
"""
Render Class Meta that contains the schema matadata
Expand Down Expand Up @@ -116,12 +116,14 @@ def render_metaclass(

if properties:
# some formating to remove identation at the end of the Class Meta to make it more compatible with black
return (
self.field_identation.join(
[line for line in templates.metaclass_template.safe_substitute(properties=properties).split("\n")]
).rstrip(self.field_identation)
+ "\n"
)
return self.field_identation.join(
[
line
for line in templates.metaclass_template.safe_substitute(
properties=properties, decorator=decorator
).split("\n")
]
).rstrip(self.field_identation)
return None

def render_docstring(self, *, docstring: typing.Optional[str]) -> str:
Expand Down Expand Up @@ -427,7 +429,7 @@ def parse_enum(self, field: JsonDict) -> str:

docstring = self.render_docstring(docstring=field.get("doc"))
enum_class = templates.enum_template.safe_substitute(name=enum_name, symbols=symbols_repr, docstring=docstring)
metaclass = self.render_metaclass(schema=field)
metaclass = self.render_metaclass(schema=field, decorator=templates.METACLASS_DECORATOR)

if metaclass:
enum_class += metaclass
Expand Down
10 changes: 8 additions & 2 deletions dataclasses_avroschema/model_generator/lang/python/templates.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
import sys
from string import Template

PYTHON_VERSION_GE_311 = sys.version_info.major == 3 and sys.version_info.minor >= 11

FIELD_TYPE_TEMPLATE = "$name: $type"
METACLASS_FIELD_TEMPLATE = '$name = "$value"'
METACLASS_ALIAS_FIELD = "$name = $value"
METACLASS_SCHEMA_FIELD = "$name = '$schema'"
METACLASS_DECORATOR = "@enum.nonmember" if PYTHON_VERSION_GE_311 else ""
FIELD_DEFAULT_TEMPLATE = " = $default"
OPTIONAL_TEMPLATE = "typing.Optional[$type]"
UNION_TEMPLATE = "typing.Union[$type]"
Expand All @@ -23,9 +27,10 @@
DECIMAL_TYPE_TEMPLATE = "types.condecimal(max_digits=$precision, decimal_places=$scale)"

ENUM_SYMBOL_TEMPLATE = "$key = $value"
ENUM_TEMPLATE = """
ENUM_PYTHON_VERSION = "str, enum.Enum" if PYTHON_VERSION_GE_311 else "enum.Enum"
ENUM_TEMPLATE = f"""
class $name(enum.Enum):$docstring
class $name({ENUM_PYTHON_VERSION}):$docstring
$symbols
"""

Expand All @@ -38,6 +43,7 @@ class $name($base_class):$docstring
INSTANCE_TEMPLATE = "$type($properties)"

METACLASS_TEMPLATE = """
$decorator
class Meta:
$properties
"""
Expand Down
190 changes: 128 additions & 62 deletions docs/case.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,68 +3,134 @@
Sometimes we use `avro schemas` with different sources (some written in Scala, some in Python, etc). With the `case`
you can generate your schemas according to your programming language convention:

```python title="Example with CAPITALCASE"
import typing
import dataclasses
import enum

from dataclasses_avroschema import AvroModel, case, types


# New enum!!
class FavoriteColor(enum.Enum):
BLUE = "BLUE"
YELLOW = "YELLOW"
GREEN = "GREEN"


@dataclasses.dataclass
class UserAdvance(AvroModel):
name: str
age: int
pets: typing.List[str]
accounts: typing.Dict[str, int]
favorite_colors: FavoriteColor
has_car: bool = False
country: str = "Argentina"
address: str = None
md5: types.Fixed = types.Fixed(16)

class Meta:
schema_doc = False


UserAdvance.avro_schema(case_type=case.CAPITALCASE)
```

resulting in

```json
{
"type": "record",
"name": "UserAdvance",
"fields": [
{"name": "Name", "type": "string"},
{"name": "Age", "type": "long"},
{"name": "Pets", "type": {
"type": "array", "items": "string", "name": "Pet"
}
},
{"name": "Accounts", "type": {
"type": "map", "values": "long", "name": "Account"
}
},
{"name": "Has_car", "type": "boolean", "default": false},
{"name": "Favorite_colors", "type": {
"type": "enum", "name": "FavoriteColor", "symbols": ["BLUE", "YELLOW", "GREEN"]
}
},
{"name": "Country", "type": "string", "default": "Argentina"},
{"name": "Address", "type": ["null", "string"], "default": null},
{"name": "Md5", "type": {"type": "fixed", "name": "Md5", "size": 16}}
]
}'
```
=== "python <= 3.10"

```python title="Example with CAPITALCASE"
import typing
import dataclasses
import enum

from dataclasses_avroschema import AvroModel, case, types


class FavoriteColor(enum.Enum):
BLUE = "BLUE"
YELLOW = "YELLOW"
GREEN = "GREEN"


@dataclasses.dataclass
class UserAdvance(AvroModel):
name: str
age: int
pets: typing.List[str]
accounts: typing.Dict[str, int]
favorite_colors: FavoriteColor
has_car: bool = False
country: str = "Argentina"
address: str = None
md5: types.Fixed = types.Fixed(16)

class Meta:
schema_doc = False


UserAdvance.avro_schema(case_type=case.CAPITALCASE)
```

resulting in

```json
{
"type": "record",
"name": "UserAdvance",
"fields": [
{"name": "Name", "type": "string"},
{"name": "Age", "type": "long"},
{"name": "Pets", "type": {
"type": "array", "items": "string", "name": "Pet"
}
},
{"name": "Accounts", "type": {
"type": "map", "values": "long", "name": "Account"
}
},
{"name": "Has_car", "type": "boolean", "default": false},
{"name": "Favorite_colors", "type": {
"type": "enum", "name": "FavoriteColor", "symbols": ["BLUE", "YELLOW", "GREEN"]
}
},
{"name": "Country", "type": "string", "default": "Argentina"},
{"name": "Address", "type": ["null", "string"], "default": null},
{"name": "Md5", "type": {"type": "fixed", "name": "Md5", "size": 16}}
]
}'
```

=== "python >= 3.11"

```python title="Example with CAPITALCASE"
import typing
import dataclasses
import enum

from dataclasses_avroschema import AvroModel, case, types


# New enum!!
class FavoriteColor(str, enum.Enum):
BLUE = "BLUE"
YELLOW = "YELLOW"
GREEN = "GREEN"


@dataclasses.dataclass
class UserAdvance(AvroModel):
name: str
age: int
pets: typing.List[str]
accounts: typing.Dict[str, int]
favorite_colors: FavoriteColor
has_car: bool = False
country: str = "Argentina"
address: str = None
md5: types.Fixed = types.Fixed(16)

class Meta:
schema_doc = False


UserAdvance.avro_schema(case_type=case.CAPITALCASE)
```

resulting in

```json
{
"type": "record",
"name": "UserAdvance",
"fields": [
{"name": "Name", "type": "string"},
{"name": "Age", "type": "long"},
{"name": "Pets", "type": {
"type": "array", "items": "string", "name": "Pet"
}
},
{"name": "Accounts", "type": {
"type": "map", "values": "long", "name": "Account"
}
},
{"name": "Has_car", "type": "boolean", "default": false},
{"name": "Favorite_colors", "type": {
"type": "enum", "name": "FavoriteColor", "symbols": ["BLUE", "YELLOW", "GREEN"]
}
},
{"name": "Country", "type": "string", "default": "Argentina"},
{"name": "Address", "type": ["null", "string"], "default": null},
{"name": "Md5", "type": {"type": "fixed", "name": "Md5", "size": 16}}
]
}'
```

*(This script is complete, it should run "as is")*

Expand Down
Loading

0 comments on commit 18e226e

Please sign in to comment.