Skip to content

RFC: use enum.StrEnum where it makes sense, when available #16603

@neutrinoceros

Description

@neutrinoceros

This idea has come up a couple times in the recent past so I figured I would start centralizing it here.
enum.StrEnum is new in Python 3.11, so we'll be able to use it when we drop support for 3.101.

There are a handful places in astropy where some exposed parameters only have a few accepted string values. It would make sense to convert user inputs to StrEnum internally to:

  1. ensure the value is acceptable (an error is raised otherwise)
  2. allow typecheckers to verify exhaustiveness in if/elif/.../else (or match/case/case/...) branches associated with these variables (this is a medium to long term goal, since we don't typecheck yet)

Since I'm using a couple relatively new concepts here that are not yet used in astropy, I don't want to assume everything I just wrote makes perfect sense to all readers, so let me illustrate my points with the following example

from enum import StrEnum, auto
from typing import Literal, assert_never

class OriginMismatch(StrEnum):
    IGNORE = auto()
	WARN = auto()
    ERROR = auto()

def _prepare_unit_sphere_coords(origin_mismatch: Literal["ignore", "warn", "error"]):
    # runtime validation implemented "for free": any unexpected value will raise a `ValueError`
    _origin_mismatch = OriginMismatch(origin_mismatch)

    match _origin_mismatch:
        case OriginMismatch.IGNORE:
            ...
        case OriginMismatch.WARN:
            ...
        case OriginMismatch.ERROR:
            ...
        case _ as unreachable:  # default
			# type checkers can verify that this branch is unreachable
			# so we cannot forget to update this block if an enum member is added or removed
            assert_never(unreachable)

This example is inspired from a real life situation, see #16246

def _prepare_unit_sphere_coords(
self,
other: BaseCoordinateFrame | SkyCoord,
origin_mismatch: Literal["ignore", "warn", "error"],
) -> tuple[Longitude, Latitude, Longitude, Latitude]:
other_frame = getattr(other, "frame", other)
if not (
origin_mismatch == "ignore"
or self.is_equivalent_frame(other_frame)
or all(
isinstance(comp, (StaticMatrixTransform, DynamicMatrixTransform))
for comp in frame_transform_graph.get_transform(
type(self), type(other_frame)
).transforms
)
):
if origin_mismatch == "warn":
warnings.warn(NonRotationTransformationWarning(self, other_frame))
elif origin_mismatch == "error":
raise NonRotationTransformationError(self, other_frame)
else:
raise ValueError(
f"{origin_mismatch=} is invalid. Allowed values are 'ignore', "
"'warn' or 'error'."
)
self_sph = self.represent_as(r.UnitSphericalRepresentation)
other_sph = other_frame.transform_to(self).represent_as(
r.UnitSphericalRepresentation
)
return self_sph.lon, self_sph.lat, other_sph.lon, other_sph.lat

Other than this example, here are other identified places where this pattern would make sense:

This list is meant to be updated (not necessarily by me) as examples come along.
Let me also ping a couple people who are already involved in this discussion: @eerovaher @nstarman @jeffjennings

Footnotes

  1. the class is actually so light weight that it can easily be backported on Python 3.10, but since this is very low priority, it might still not be worth doing.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions