-
Notifications
You must be signed in to change notification settings - Fork 152
Expand file tree
/
Copy pathgenerated_enum.py
More file actions
85 lines (64 loc) · 1.97 KB
/
generated_enum.py
File metadata and controls
85 lines (64 loc) · 1.97 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import re
from dataclasses import dataclass
from typing import Optional
import codegen.packages as packages
from codegen.code_builder import CodeBuilder
from codegen.generated_dataclass import _append_description
from codegen.jsonschema import Schema, Stage
@dataclass(kw_only=True)
class GeneratedEnum:
class_name: str
package: str
values: dict[str, str]
description: Optional[str]
experimental: bool
deprecated: bool
def generate_enum(namespace: str, schema_name: str, schema: Schema) -> GeneratedEnum:
assert schema.enum
class_name = packages.get_class_name(schema_name)
package = packages.get_package(namespace, schema_name)
values = {}
assert package
for value in schema.enum:
values[_camel_to_upper_snake(value)] = value
return GeneratedEnum(
class_name=class_name,
package=package,
values=values,
description=schema.description,
experimental=schema.stage == Stage.PRIVATE,
deprecated=schema.deprecated or False,
)
def get_code(generated: GeneratedEnum) -> str:
b = CodeBuilder()
# Example:
#
# class Color(Enum):
#
b.append(f"class {generated.class_name}(Enum):")
b.newline()
_append_description(
b,
generated.description,
experimental=generated.experimental,
deprecated=generated.deprecated,
)
# Example:
#
# RED = "RED"
#
for key, value in generated.values.items():
b.indent().append(f'{key} = "{value}"')
b.newline()
b.newline()
# Example:
#
# ColorParam = Literal["RED", "GREEN", "BLUE"] | Color
b.append(generated.class_name).append('Param = Literal["')
b.append_list(list(generated.values.values()), sep='", "')
b.append('"] | ', generated.class_name)
b.newline()
return b.build()
def _camel_to_upper_snake(value):
s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", value)
return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).upper()