Enum Generator: allow passing strings for parameters which are enums to make calling more pythonic

[grusin-db]

Problem

Generated Enum types does not have support of being instantiated from string, hence it forces end users to not use non pythonic syntax of:

from databricks.sdk.service.workspace import ExportFormat

w.workspace.export('/my_notebook', format=ExportFormat.SOURCE)

instead of neat, short, and pythonic:

w.workspace.export('/my_notebook', format='SOURCE')

throws an exception:

AttributeError                            Traceback (most recent call last)
<command-3538949529976490> in <cell line: 1>()
----> 1 w.workspace.export('/Users/[email protected]/auto-dlt/generated/gr_products/gr_bronze-autoloader', format='SOURCE')

/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/databricks/sdk/service/workspace.py in export(self, path, direct_download, format, **kwargs)
    225         query = {}
    226         if direct_download: query['direct_download'] = request.direct_download
--> 227         if format: query['format'] = request.format.value
    228         if path: query['path'] = request.path
    229 

AttributeError: 'str' object has no attribute 'value'

Example

For example:

help(w.workspace.export)

has signature of:

export(path: str, *, direct_download: bool = None, format: databricks.sdk.service.workspace.ExportFormat = None, **kwargs) -> databricks.sdk.service.workspace.ExportResponse method of databricks.sdk.service.workspace.WorkspaceAPI instance

in turn parameter format is of Enum type defined in databricks.sdk.service.workspace.ExportFormat as

class ExportFormat(Enum):
    """This specifies the format of the file to be imported. By default, this is `SOURCE`. However it
    may be one of: `SOURCE`, `HTML`, `JUPYTER`, `DBC`. The value is case sensitive."""

    DBC = 'DBC'
    HTML = 'HTML'
    JUPYTER = 'JUPYTER'
    R_MARKDOWN = 'R_MARKDOWN'
    SOURCE = 'SOURCE'

[grusin-db]

I am not sure what is the neat/correct way of solving it, SO has all sort of answers (https://stackoverflow.com/questions/41407414/convert-string-to-enum-in-python)

this stands out as a neat one: https://stackoverflow.com/a/56567247

[grusin-db]

from enum import Enum

class YetAnotherStrEnum(Enum):
  @classmethod
  def from_str(cls, s: str):
    return cls.__members__.get(s)

  @classmethod
  def from_enum_or_str(cls, s_or_e):
    if isinstance(s_or_e, str):
      return cls.from_str(s_or_e)
    else:
      return s_or_e

class ExportFormat(YetAnotherStrEnum):
    """This specifies the format of the file to be imported. By default, this is `SOURCE`. However it
    may be one of: `SOURCE`, `HTML`, `JUPYTER`, `DBC`. The value is case sensitive."""

    DBC = 'DBC'
    HTML = 'HTML'
    JUPYTER = 'JUPYTER'
    R_MARKDOWN = 'R_MARKDOWN'
    SOURCE = 'SOURCE'
    aaa = 'bbb'

assert ExportFormat.from_str('SOURCE') == ExportFormat.SOURCE
assert ExportFormat.from_enum_or_str('SOURCE') == ExportFormat.SOURCE
assert ExportFormat.from_enum_or_str(ExportFormat.SOURCE) == ExportFormat.SOURCE

hence calling code can be change from:

if format: query['format'] = request.format.value

to:

if format: query['format'] = ExportFormat.from_enum_or_str(request.format).value

this way both ExportFormat Enum and str would be handled.

[grusin-db]

Another issue

...with displaying/repr for classes having enums, displayed text has nasty looking object_type and language, making copy & paste unfriendly

ObjectInfo(created_at=1677658169185, language=<Language.SQL: 'SQL'>, modified_at=1677658409090, object_id=3538949529947478, object_type=<ObjectType.NOTEBOOK: 'NOTEBOOK'>, path='/Users/[email protected]/auto-dlt/generated/gr_products/gr_silver-incremental', size=None)]

displayed text should look like this, so it can be copied and run to created object back:

from databricks.sdk.service.workspace import ObjectType, Language, ObjectInfo

ObjectInfo(created_at=1677658169185, language=Language.SQL, modified_at=1677658409090, object_id=3538949529947478, object_type=ObjectType.NOTEBOOK, path='/Users/[email protected]/auto-dlt/generated/gr_products/gr_silver-incremental', size=None)

Solution:

adding custom string formatter:

class YetAnotherStrEnum(Enum):
  ...

  def __repr__(self):
    return str(self)

gets rid of nasty EnumType.FieldName: FieldValue formatting, in favour of just EnumType.FieldName