cli/python/codegen/codegen/generated_enum.py at main · databricks/cli · GitHub

History

85 lines (64 loc) · 1.99 KB

Raw

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 LaunchStage, Schema

@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 == LaunchStage.PRIVATE_PREVIEW,

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(

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()

# 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()

cli/python/codegen/codegen/generated_enum.py at main · databricks/cli · GitHub

Uh oh!

Expand file tree

Breadcrumbs

generated_enum.py

Latest commit

History

Breadcrumbs

generated_enum.py

File metadata and controls