Feat: --known-enum-bases option to support custom enum-like base cl… (#3513)

Author: imtoopunkforyou

CHANGELOG.md

@@ -23,6 +23,7 @@ Semantic versioning in our case means:
 
 - Allows `__init__.py` files that consist only of imports, #3486
 - Adds `--max-conditions` option, #3493
+- Adds `--known-enum-bases` option to support custom enum-like base classes, #3513
 
 ### Misc
 

tests/test_visitors/test_ast/test_naming/test_class_attributes.py

@@ -44,6 +44,48 @@ class Test({0}):
     'lower, UPPER = field',
 ]
 
+custom_enum_base_class = """
+class BaseEnum(enum.Enum):
+    def __init__(self, short: str, long: str) -> None:
+        self.short = short
+        self.long = long
+
+class ConcreteEnum(BaseEnum):
+    HELLO = 'h', 'hello'
+    WORLD = 'w', 'world'
+"""
+
+
+multiple_custom_enum_bases = """
+class FirstEnum(enum.Enum):
+    pass
+
+class SecondEnum(enum.Enum):
+    pass
+
+class FirstEnumSuccessor(FirstEnum):
+    HELLO = 'hello'
+    WORLD = 'world'
+
+class SecondEnumSuccessor(SecondEnum):
+    PING = 'pong'
+
+class MultiInheritanceEnum(FirstEnum, SecondEnum):
+    MULTI = 'multi'
+"""
+
+dotted_names_enum_bases = """
+from src import constants
+
+class FirstEnum(constants.MyEnum):
+    HELLO = 'hello'
+    WORLD = 'world'
+
+class SecondEnum(constants.MyEnum):
+    HELLO = 'hello'
+    WORLD = 'world'
+"""
+
 
 @pytest.mark.parametrize(
     'code',
@@ -202,3 +244,56 @@ def test_regression423(
     visitor.run()
 
     assert_errors(visitor, [])
+
+
+@pytest.mark.parametrize('with_configuration', [True, False])
+def test_custom_enum_base_class_with_config(
+    assert_errors,
+    parse_ast_tree,
+    options,
+    with_configuration,
+):
+    """Testing that custom enum base classes work with configuration."""
+    tree = parse_ast_tree(custom_enum_base_class)
+
+    if with_configuration:
+        options_with_config = options(known_enum_bases=('BaseEnum',))
+        visitor = WrongNameVisitor(options_with_config, tree=tree)
+        visitor.run()
+        assert_errors(visitor, [])
+    else:
+        options_without_config = options()
+        visitor = WrongNameVisitor(options_without_config, tree=tree)
+        visitor.run()
+        assert_errors(
+            visitor,
+            [UpperCaseAttributeViolation, UpperCaseAttributeViolation],
+        )
+
+
+def test_multiple_custom_enum_bases_with_config(
+    assert_errors,
+    parse_ast_tree,
+    options,
+):
+    """Testing that multiple custom enum base classes work with config."""
+    tree = parse_ast_tree(multiple_custom_enum_bases)
+
+    options = options(known_enum_bases=('FirstEnum', 'SecondEnum'))
+    visitor = WrongNameVisitor(options, tree=tree)
+    visitor.run()
+    assert_errors(visitor, [])
+
+
+def test_dotted_names_in_enum_bases_config(
+    assert_errors,
+    parse_ast_tree,
+    options,
+):
+    """Testing that dotted names in enum bases config work correctly."""
+    tree = parse_ast_tree(dotted_names_enum_bases)
+
+    options = options(known_enum_bases=('constants.MyEnum',))
+    visitor = WrongNameVisitor(options, tree=tree)
+    visitor.run()
+    assert_errors(visitor, [])

wemake_python_styleguide/logic/naming/enums.py

@@ -3,6 +3,7 @@
 from typing import Final
 
 from wemake_python_styleguide.logic.source import node_to_string
+from wemake_python_styleguide.options.validation import ValidatedOptions
 
 _CONCRETE_ENUM_NAMES: Final = (
     'enum.StrEnum',
@@ -68,3 +69,26 @@ def has_enum_like_base(defn: ast.ClassDef) -> bool:
     https://docs.djangoproject.com/en/5.1/ref/models/fields/#choices
     """
     return _has_one_of_base_classes(defn, _ENUM_LIKE_NAMES)
+
+
+def has_enum_like_base_with_config(
+    defn: ast.ClassDef,
+    config: ValidatedOptions,
+) -> bool:
+    """
+    Tells if some class has `Enum` or semantically similar class as its base.
+
+    This function also checks user-defined enum-like base classes from
+    configuration.
+    """
+    if has_enum_like_base(defn):
+        return True
+
+    enum_bases = config.known_enum_bases
+    if enum_bases:
+        normalized: tuple[str, ...] = tuple({
+            name for base in enum_bases for name in (base, base.split('.')[-1])
+        })
+        return _has_one_of_base_classes(defn, normalized)
+
+    return False

wemake_python_styleguide/options/config.py

@@ -58,6 +58,9 @@ class of violations that are forbidden to ignore inline, defaults to
 - ``exps-for-one-empty-line`` - number of expressions in
     function or method body for available empty line (without code)
     :str:`wemake_python_styleguide.options.defaults.EXPS_FOR_ONE_EMPTY_LINE`
+- ``known-enum-bases`` - list of additional enum-like base class names that
+    should be treated as enums, defaults to
+    :str:`wemake_python_styleguide.options.defaults.KNOWN_ENUM_BASES`
 
 
 .. rubric:: Complexity options
@@ -283,6 +286,14 @@ class Configuration:
             defaults.EXPS_FOR_ONE_EMPTY_LINE,
             'Count of expressions for one empty line in a function body.',
         ),
+        _Option(
+            '--known-enum-bases',
+            defaults.KNOWN_ENUM_BASES,
+            'List of additional enum-like base class names that should be '
+            'treated as enums.',
+            type=String,
+            comma_separated_list=True,
+        ),
         # Complexity:
         _Option(
             '--max-returns',

wemake_python_styleguide/options/defaults.py

@@ -34,6 +34,9 @@
     'Config',  # pydantic specific
 )
 
+#: List of additional enum-like base class names.
+KNOWN_ENUM_BASES: Final = ()
+
 #: Domain names that are removed from variable names' blacklist.
 ALLOWED_DOMAIN_NAMES: Final = ()
 

wemake_python_styleguide/options/validation.py

@@ -63,6 +63,7 @@ class ValidatedOptions:
         validator=[_min_max(min=1, max=defaults.MAX_NOQA_COMMENTS)],
     )
     nested_classes_whitelist: tuple[str, ...] = attr.ib(converter=tuple)
+    known_enum_bases: tuple[str, ...] = attr.ib(converter=tuple)
     allowed_domain_names: tuple[str, ...] = attr.ib(converter=tuple)
     forbidden_domain_names: tuple[str, ...] = attr.ib(converter=tuple)
     allowed_module_metadata: tuple[str, ...] = attr.ib(converter=tuple)

wemake_python_styleguide/violations/naming.py

@@ -448,6 +448,11 @@ class UpperCaseAttributeViolation(ASTViolation):
         Rename your variables so that they conform
         to ``snake_case`` convention.
 
+    Configuration:
+        This rule is configurable with ``--known-enum-bases``.
+        Default:
+        :str:`wemake_python_styleguide.options.defaults.KNOWN_ENUM_BASES`
+
     Example::
 
         # Correct:

wemake_python_styleguide/visitors/ast/naming/validation.py

@@ -228,7 +228,7 @@ def check_attribute_names(self, node: ast.ClassDef) -> None:
             node,
             include_annotated=True,
         )
-        is_enum_like = enums.has_enum_like_base(node)
+        is_enum_like = enums.has_enum_like_base_with_config(node, self._options)
 
         for assign in class_attributes:
             for target in get_assign_targets(assign):