Skip to content

Commit

Permalink
use enum suffix setting to generate enum names
Browse files Browse the repository at this point in the history
  • Loading branch information
jrbeaumont committed Jan 5, 2024
1 parent fbc9ca8 commit 17f0888
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 6 deletions.
15 changes: 9 additions & 6 deletions drf_spectacular/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,9 @@ def extract_hash(schema):
prop_hash_mapping[prop_name].add(prop_enum_cleaned_hash)
hash_name_mapping[prop_enum_cleaned_hash].add((component_name, prop_name))

# get the suffix to be used for enums from settings
enum_suffix = spectacular_settings.ENUM_SUFFIX

# traverse all enum properties and generate a name for the choice set. naming collisions
# are resolved and a warning is emitted. giving a choice set multiple names is technically
# correct but potentially unwanted. also emit a warning there to make the user aware.
Expand All @@ -84,13 +87,13 @@ def extract_hash(schema):
enum_name = overrides[prop_hash]
elif len(prop_hash_set) == 1:
# prop_name has been used exclusively for one choice set (best case)
enum_name = f'{camelize(prop_name)}Enum'
enum_name = f'{camelize(prop_name)}{enum_suffix}'
elif len(hash_name_mapping[prop_hash]) == 1:
# prop_name has multiple choice sets, but each one limited to one component only
component_name, _ = next(iter(hash_name_mapping[prop_hash]))
enum_name = f'{camelize(component_name)}{camelize(prop_name)}Enum'
enum_name = f'{camelize(component_name)}{camelize(prop_name)}{enum_suffix}'
else:
enum_name = f'{camelize(prop_name)}{prop_hash[:3].capitalize()}Enum'
enum_name = f'{camelize(prop_name)}{prop_hash[:3].capitalize()}{enum_suffix}'
warn(
f'enum naming encountered a non-optimally resolvable collision for fields '
f'named "{prop_name}". The same name has been used for multiple choice sets '
Expand Down Expand Up @@ -143,12 +146,12 @@ def extract_hash(schema):
]
if spectacular_settings.ENUM_ADD_EXPLICIT_BLANK_NULL_CHOICE:
if '' in prop_enum_original_list:
components.append(create_enum_component('BlankEnum', schema={'enum': ['']}))
components.append(create_enum_component(f'Blank{enum_suffix}', schema={'enum': ['']}))
if None in prop_enum_original_list:
if spectacular_settings.OAS_VERSION.startswith('3.1'):
components.append(create_enum_component('NullEnum', schema={'type': 'null'}))
components.append(create_enum_component(f'Null{enum_suffix}', schema={'type': 'null'}))
else:
components.append(create_enum_component('NullEnum', schema={'enum': [None]}))
components.append(create_enum_component(f'Null{enum_suffix}', schema={'enum': [None]}))

# undo OAS 3.1 type list NULL construction as we cover this in a separate component already
if spectacular_settings.OAS_VERSION.startswith('3.1') and isinstance(enum_schema['type'], list):
Expand Down
19 changes: 19 additions & 0 deletions tests/test_postprocessing.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,3 +389,22 @@ def get(self, request):
assert schema['components']['schemas']['SomeTestEnum'] == {
'enum': [0, 1], 'type': 'integer', 'description': '* `0` - test group A\n* `1` - test group B',
}


@mock.patch('drf_spectacular.settings.spectacular_settings.ENUM_NAME_OVERRIDES', {
'VoteChoices': 'tests.test_postprocessing.vote_choices'
})
def test_enum_suffix(no_warnings, clear_caches):
"""Test that enums generated have the suffix from the settings."""
# check variations of suffix
enum_suffix_variations = ['Type', 'Enum', 'Testing', '']
for variation in enum_suffix_variations:
with mock.patch('drf_spectacular.settings.spectacular_settings.ENUM_SUFFIX', variation):
schema = generate_schema('a', AViewset)

assert f'Null{variation}' in schema['components']['schemas']
assert f'Blank{variation}' in schema['components']['schemas']
assert f'Language{variation}' in schema['components']['schemas']
# vote choices is overridden, so should not have the suffix added
assert f'Vote{variation}' not in schema['components']['schemas']
assert 'VoteChoices' in schema['components']['schemas']

0 comments on commit 17f0888

Please sign in to comment.