diff --git a/drf_spectacular/hooks.py b/drf_spectacular/hooks.py index df9d30df..c4d2f8a7 100644 --- a/drf_spectacular/hooks.py +++ b/drf_spectacular/hooks.py @@ -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. @@ -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 ' @@ -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): diff --git a/drf_spectacular/settings.py b/drf_spectacular/settings.py index fc29b95f..dae6330e 100644 --- a/drf_spectacular/settings.py +++ b/drf_spectacular/settings.py @@ -120,6 +120,9 @@ 'ENUM_ADD_EXPLICIT_BLANK_NULL_CHOICE': True, # Add/Append a list of (``choice value`` - choice name) to the enum description string. 'ENUM_GENERATE_CHOICE_DESCRIPTION': True, + # Optional suffix for generated enum. + # e.g. {'ENUM_SUFFIX': "Type"} would produce an enum name 'StatusType'. + 'ENUM_SUFFIX': 'Enum', # function that returns a list of all classes that should be excluded from doc string extraction 'GET_LIB_DOC_EXCLUDES': 'drf_spectacular.plumbing.get_lib_doc_excludes', diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index f671a3b2..1f08e06d 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -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']