A pytest example to see if the Enum members are sorted by value

README.md

Sometimes you just want the StrEnums in python to be sorted by their values to avoid duplicates and also to ensure that the code is maintainable.

This test takes the selected files and then runs a test using AST to see if the enum members are sorted by the member values.

If not the test fails.

enum-value-sorting-test-.py

import ast

import pytest

filenames = ["integrations/edi/enums.py"]


class EnumVisitor(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.errors = []

    def visit_ClassDef(self, node):
        if (
            node.bases
            and isinstance(node.bases[0], ast.Name)
            and node.bases[0].id == "StrEnum"
        ):
            # Get all the assignments, which are the Enum members
            members = [b.value for b in node.body if isinstance(b, ast.Assign)]

            # Get the values of the assignments
            values = [member.value for member in members]
            if values != sorted(values):
                self.errors.append(
                    f"Enum '{node.name}' is not sorted by value in line {node.lineno}"
                )
                return


def check_file(filename):
    with open(filename, "r") as source:
        tree = ast.parse(source.read())
        visitor = EnumVisitor()
        visitor.visit(tree)
        return visitor.errors


@pytest.mark.parametrize("filename", filenames)
def test_enum_sorting(filename):
    if errors := check_file(filename):
        pytest.fail("\n".join(errors))