diff --git a/omymodels/models/enum/template.jinja2 b/omymodels/models/enum/template.jinja2 index fb4bcda..ef93a95 100644 --- a/omymodels/models/enum/template.jinja2 +++ b/omymodels/models/enum/template.jinja2 @@ -1,6 +1,10 @@ {% for _type in custom_types %} - -class {{_type['name']}}({{_type['parents']}}): - - {% for key, value in _type['properties']['values'].items() %}{{key}} = {{value}} - {% endfor %}{% endfor %} \ No newline at end of file +{{ _type['name'] }} = Enum( + value='{{ _type['name'] }}', + names=[ + {%- for key, value in _type['properties']['values'].items() %} + ('{{ key }}', {{ value }}){% if not loop.last %},{% endif %} + {%- endfor %} + ] +) +{% endfor %} \ No newline at end of file diff --git a/tests/functional/converter/test_converter.py b/tests/functional/converter/test_converter.py index 5ee3508..1c2d182 100644 --- a/tests/functional/converter/test_converter.py +++ b/tests/functional/converter/test_converter.py @@ -34,12 +34,14 @@ class Material: db = Gino() +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', article), + ('video', video) + ] +) -class MaterialType(str, Enum): - - article = article - video = video - class Material(db.Model): @@ -100,12 +102,14 @@ def test_from_sqlalchemy_to_gino(): db = Gino() +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(db.Model): @@ -161,12 +165,14 @@ def test_from_sqlalchemy_to_pydantic(): from typing import Optional from pydantic import BaseModel, Json +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(BaseModel): diff --git a/tests/functional/generator/test_dataclasses.py b/tests/functional/generator/test_dataclasses.py index 87f42b9..5c530d1 100644 --- a/tests/functional/generator/test_dataclasses.py +++ b/tests/functional/generator/test_dataclasses.py @@ -73,12 +73,14 @@ def test_enums_in_dataclasses(): from typing import Union from dataclasses import dataclass, field +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - @dataclass class Material: @@ -119,12 +121,14 @@ def test_defaults_off(): from typing import Union from dataclasses import dataclass, field +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - @dataclass class Material: @@ -165,12 +169,14 @@ def test_upper_now_produces_same_result(): from typing import Union from dataclasses import dataclass, field +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - @dataclass class Material: diff --git a/tests/functional/generator/test_enum_only.py b/tests/functional/generator/test_enum_only.py index e375a06..dc1ed3d 100644 --- a/tests/functional/generator/test_enum_only.py +++ b/tests/functional/generator/test_enum_only.py @@ -12,12 +12,14 @@ def test_enum_only(): expected = """from enum import Enum - -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - """ +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) +""" assert result == expected @@ -33,18 +35,42 @@ def test_enum_models(): result = create_models(ddl)["code"] expected = """from enum import Enum, IntEnum +ContentType = Enum( + value='ContentType', + names=[ + ('HTML', 'HTML'), + ('MARKDOWN', 'MARKDOWN'), + ('TEXT', 'TEXT') + ] +) -class ContentType(str, Enum): - - HTML = 'HTML' - MARKDOWN = 'MARKDOWN' - TEXT = 'TEXT' - - -class Period(IntEnum): +Period = Enum( + value='Period', + names=[ + ('zero', 0), + ('one', 1), + ('two', 2) + ] +) +""" + assert expected == result - zero = 0 - one = 1 - two = 2 +def test_enum_with_whitespace(): + ddl = """ + CREATE TYPE "enum_with_spaces" AS ENUM ( + 'some value', + 'another value' + ); """ - assert expected == result + result = create_models(ddl)["code"] + expected = """from enum import Enum + +EnumWithSpaces = Enum( + value='EnumWithSpaces', + names=[ + ('another value', 'another value'), + ('some value', 'some value') + ] +) +""" + assert result == expected \ No newline at end of file diff --git a/tests/functional/generator/test_enums.py b/tests/functional/generator/test_enums.py index 2b9c711..122fe4d 100644 --- a/tests/functional/generator/test_enums.py +++ b/tests/functional/generator/test_enums.py @@ -7,13 +7,15 @@ def test_enums_gino(): db = Gino(schema="schema--notification") +ContentType = Enum( + value='ContentType', + names=[ + ('HTML', 'HTML'), + ('MARKDOWN', 'MARKDOWN'), + ('TEXT', 'TEXT') + ] +) -class ContentType(str, Enum): - - HTML = 'HTML' - MARKDOWN = 'MARKDOWN' - TEXT = 'TEXT' - class Notification(db.Model): @@ -51,20 +53,24 @@ def test_pydantic_models(): from typing import Optional from pydantic import BaseModel +ContentType = Enum( + value='ContentType', + names=[ + ('HTML', 'HTML'), + ('MARKDOWN', 'MARKDOWN'), + ('TEXT', 'TEXT') + ] +) + +Period = Enum( + value='Period', + names=[ + ('zero', 0), + ('one', 1), + ('two', 2) + ] +) -class ContentType(str, Enum): - - HTML = 'HTML' - MARKDOWN = 'MARKDOWN' - TEXT = 'TEXT' - - -class Period(IntEnum): - - zero = 0 - one = 1 - two = 2 - class Notification(BaseModel): @@ -102,12 +108,14 @@ def test_enum_works_with_lower_case(): db = Gino() +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(db.Model): diff --git a/tests/functional/generator/test_pydantic_models.py b/tests/functional/generator/test_pydantic_models.py index ede309c..a9084c6 100644 --- a/tests/functional/generator/test_pydantic_models.py +++ b/tests/functional/generator/test_pydantic_models.py @@ -92,12 +92,14 @@ def test_enums_lower_case_names_works(): from typing import Optional from pydantic import BaseModel, Json +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(BaseModel): @@ -157,12 +159,14 @@ def test_no_defaults(): from typing import Optional from pydantic import BaseModel, Json +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(BaseModel): diff --git a/tests/functional/generator/test_sqlalchemy.py b/tests/functional/generator/test_sqlalchemy.py index 118d332..523ecfe 100644 --- a/tests/functional/generator/test_sqlalchemy.py +++ b/tests/functional/generator/test_sqlalchemy.py @@ -11,12 +11,14 @@ def test_with_enums(): Base = declarative_base() +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(Base): @@ -263,12 +265,14 @@ def test_upper_name_produces_the_same_result(): Base = declarative_base() +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(Base): diff --git a/tests/functional/generator/test_sqlmodel.py b/tests/functional/generator/test_sqlmodel.py index 32b0e48..50f3a18 100644 --- a/tests/functional/generator/test_sqlmodel.py +++ b/tests/functional/generator/test_sqlmodel.py @@ -12,12 +12,14 @@ def test_with_enums(): from sqlalchemy.dialects.postgresql import JSON from pydantic import Json, UUID4 +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(SQLModel, table=True): @@ -262,12 +264,14 @@ def test_upper_name_produces_the_same_result(): from sqlalchemy.dialects.postgresql import JSON from pydantic import Json, UUID4 +MaterialType = Enum( + value='MaterialType', + names=[ + ('article', 'article'), + ('video', 'video') + ] +) -class MaterialType(str, Enum): - - article = 'article' - video = 'video' - class Material(SQLModel, table=True):