python - 枚举所有可能的数据类实例(仅包含枚举和 bool 字段)

Question

基本上我需要跟随。我有一个python3 dataclass或NamedTuple ，仅包含 enum 和 bool 字段。例如:

from enum import Enum, auto
from typing import NamedTuple

class MyEnum(Enum):
    v1 = auto()
    v2 = auto()
    v3 = auto()


class MyStateDefinition(NamedTuple):
    a: MyEnum
    b: bool

这里是否有任何众所周知的解决方案来枚举此类数据类的所有可能的不相等实例？ (上面的示例有 6 个可能的不相等实例)。

也许它不是我应该使用的数据类，而是其他东西。或者我应该玩像 dataclasses.fields 这样的东西直接？

我把它想象成一些表生成器，它接受namedtuple或dataclass作为输入参数并生成所有可能的值。

table = DataTable(MyStateDefinition)
for item in table:
    # Use items somehow
    print(item.a)
    print(item.b)

为什么我需要它？我只有一些由枚举和 bool 组成的状态定义。我相信它可以作为位掩码来实现。但当涉及到用新值扩展位掩码时，结果就是一场噩梦。毕竟，位掩码似乎是一种非 Pythonic 的做事方式。

目前我必须使用我自己的实现。但也许我正在重新发明轮子。

谢谢!

Answer 1

您可以使用枚举来完成此操作，并将数据元组作为枚举成员的值(如果您愿意，可以使用 Enum/NamedTuple 混合体)。 _ignore_ 属性用于防止类命名空间中的某些名称被转换为枚举成员。

from itertools import product
from enum import Enum

class Data(Enum):
    _ignore_ = "Data", "myenum_member", "truthiness"

    @property
    def a(self):
        return self.value[0]

    @property
    def b(self):
        return self.value[1]

    def __repr__(self):
        return f'Data(a={self.a!r}, b={self.b!r})'

    Data = vars()
    for myenum_member, truthiness in product(MyEnum, (True, False)):
        Data[f'{myenum_member.name}_{truthiness}'] = (myenum_member, truthiness)

您应该能够按照您的意愿迭代生成的枚举类。

枚举的这种用法类似于 "time period" example在文档的 Enum HOWTO 部分。

---

动态生成此类表格

如果你想动态生成这种表，你可以这样做，(ab)使用元类。我已经展示了如何在文档字符串中使用此 DataTable 类的示例用法。 (出于某种原因，在 doctest 中使用 typing.get_type_hints 似乎会导致 doctest 模块出错，但如果您自己在交互式终端中尝试，这些示例确实有效。)我决定使用特殊情况 typing.Literal ，而不是像您在回答中所做的那样使用特殊情况 bool ，因为它似乎是一个更具可扩展性的选项(并且bool 可以拼写为 typing.Literal[True, False])。

from __future__ import annotations
from itertools import product
from enum import Enum, EnumMeta

from typing import (
    Iterable,
    Mapping,
    cast,
    Protocol,
    get_type_hints,
    Any,
    get_args,
    get_origin,
    Literal,
    TypeVar,
    Union,
    Optional
)

D = TypeVar('D')
T = TypeVar('T')


class DataTableFactory(EnumMeta):
    """A helper class for making data tables (an implementation detail of `DataTable`)."""

    _CLS_BASES = (Enum,)

    @classmethod
    def __prepare__(  # type: ignore[override]
            metacls,
            cls_name: str,
            fields: Mapping[str, Iterable[Any]]
    ) -> dict[str, Any]:

        cls_dict = cast(
            dict[str, Any],
            super().__prepare__(cls_name, metacls._CLS_BASES)
        )

        for i, field in enumerate(fields.keys()):
            cls_dict[field] = property(fget=lambda self, i=i: self.value[i])  # type: ignore[misc]

        for p in product(*fields.values()):
            cls_dict['_'.join(map(str, p))] = p

        def __repr__(self: Enum) -> str:
            contents = ', '.join(
                f'{field}={getattr(self, field)!r}'
                for field in fields
            )
            return f'{cls_name}Member({contents})'

        cls_dict['__repr__'] = __repr__
        return cls_dict

    @classmethod
    def make_datatable(
            metacls,
            cls_name: str,
            *,
            fields: Mapping[str, Iterable[Any]],
            doc: Optional[str] = None
    ) -> type[Enum]:
        """Create a new data table"""

        cls_dict = metacls.__prepare__(cls_name, fields)
        new_cls = metacls.__new__(metacls, cls_name, metacls._CLS_BASES, cls_dict)
        new_cls.__module__ = __name__

        if doc is None:
            all_attrs = '\n'.join(
                f'    {f"{attr_name}: ":<{(max(map(len, fields)) + 3)}}one of {attr_val!r}'
                for attr_name, attr_val in fields.items()
            )

            fields_len = len(fields)

            doc = (
                f'An enum-like data table.\n\n'
                f'All members of this data table have {fields_len} '
                f'read-only attribute{"s" if fields_len > 1 else ""}:\n'
                f'{all_attrs}\n\n'
                f'----------------------------------------------------------------------'
            )

        new_cls.__doc__ = doc
        return cast(type[Enum], new_cls)

    def __repr__(cls) -> str:
        return f"<Data table '{cls.__name__}'>"

    def index_of(cls: Iterable[D], member: D) -> int:
        """Get the index of a member in the list of members."""
        return list(cls).index(member)

    def get(
            cls: Iterable[D],
            /,
            *,
            default_: Optional[T] = None,
            **kwargs: Any
    ) -> Union[D, T, None]:
        """Return instance for given arguments set.
        Return `default_` if no member matches those arguments.
        """

        it = (
            member for member in cls
            if all((getattr(member, key) == val) for key, val in kwargs.items())
        )

        return next(it, default_)

    def __dir__(cls) -> list[str]:
        # By defining __dir__, we make methods defined in this class
        # discoverable by the interactive help() function in the REPL
        return cast(list[str], super().__dir__()) + ['index_of', 'get']


class TypedStructProto(Protocol):
    """In order to satisfy this interface, a type must have an __annotations__ dict."""
    __annotations__: dict[str, Union[Iterable[Any], type[Literal[True]]]]


class DataTableMeta(type):
    """Metaclass for `DataTable`."""
    __call__ = DataTableFactory.make_datatable  # type: ignore[assignment]


class DataTable(metaclass=DataTableMeta):
    """A mechanism to create 'data table enumerations' -- not really a class at all!

    Example usage
    -------------
    >>> Cars = DataTable('Cars', fields={'make': ('Toyota', 'Audi'), 'colour': ('Red', 'Blue')})
    >>> Cars
    <Data table 'Cars'>
    >>> list(Cars)
    [CarsMember(make=Toyota, colour=Red), CarsMember(make=Toyota, colour=Blue), CarsMember(make=Audi, colour=Red), CarsMember(make=Audi, colour=Blue)]
    >>> Cars.get(make='Audi', colour='Red')
    CarsMember(make=Audi, colour=Red)
    >>> Cars.index_of(_)
    2
    """

    @classmethod
    def from_struct(cls, cls_name: str, *, struct: type[TypedStructProto], doc: Optional[str] = None) -> type[Enum]:
        """Make a DataTable from a "typed struct" -- e.g. a dataclass, NamedTuple or TypedDict.

        Example usage (works the same way with dataclasses and TypedDicts)
        -------------------------------------------------------------------
        >>> from enum import Enum, auto
        >>> from typing import NamedTuple, Literal
        >>> class E(Enum):
        ...     v1 = auto()
        ...     v2 = auto()
        ...     v3 = auto()
        ...
        >>> class BoolsEndEnums(NamedTuple):
        ...     a: E
        ...     b: Literal[True, False]
        ...
        >>> BoolsEndEnumsTable = DataTable.from_struct('BoolsEndEnumsTable', struct=BoolsEndEnums)
        >>> list(BoolsEndEnumsTable)
        [BoolsEndEnumsTableMember(a=E.v1, b=True), BoolsEndEnumsTableMember(a=E.v1, b=False), BoolsEndEnumsTableMember(a=E.v2, b=True), BoolsEndEnumsTableMember(a=E.v2, b=False), BoolsEndEnumsTableMember(a=E.v3, b=True), BoolsEndEnumsTableMember(a=E.v3, b=False)]
        """

        fields = get_type_hints(struct)

        for field_name, field_val in fields.items():
            if get_origin(field_val) is Literal:
                fields[field_name] = get_args(field_val)

        return cast(type[Enum], cls(cls_name, fields=fields, doc=doc))  # type: ignore[call-arg]

我必须使用类型提示做一些“有趣”的事情，但 MyPy 是 sort of happy有了这一切。