-
-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Description
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:
- ensure the value is acceptable (an error is raised otherwise)
- allow typecheckers to verify exhaustiveness in
if/elif/.../else(ormatch/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
astropy/astropy/coordinates/baseframe.py
Lines 1689 to 1718 in cacde01
| 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:
- TYP: add type annotations to a couple utility functions in nddata #16546 (comment)
- Add type annotations to
statssubpackage #16562 (comment) astropy.timeseries'slombscargle'snormalizationargument could be validated early as an enum
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
-
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. ↩