|
| 1 | +"""Generic operations derived directly from Python type annotations.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import types |
| 6 | +from enum import Enum |
| 7 | +from typing import Annotated, Union, get_args, get_origin |
| 8 | + |
| 9 | + |
| 10 | +def is_union_type(annotation: object) -> bool: |
| 11 | + """Return whether ``annotation`` is a typing or PEP 604 union.""" |
| 12 | + |
| 13 | + return get_origin(annotation) in {Union, types.UnionType} |
| 14 | + |
| 15 | + |
| 16 | +def optional_member_type(annotation: object) -> object | None: |
| 17 | + """Return ``T`` only for the simple optional declaration ``T | None``.""" |
| 18 | + |
| 19 | + if not is_union_type(annotation): |
| 20 | + return None |
| 21 | + members = get_args(annotation) |
| 22 | + if len(members) != 2 or type(None) not in members: |
| 23 | + return None |
| 24 | + return next(member for member in members if member is not type(None)) |
| 25 | + |
| 26 | + |
| 27 | +def make_optional(annotation: object) -> object: |
| 28 | + """Return the same declaration with ``None`` admitted exactly once.""" |
| 29 | + |
| 30 | + if is_union_type(annotation) and type(None) in get_args(annotation): |
| 31 | + return annotation |
| 32 | + return Union[annotation, type(None)] |
| 33 | + |
| 34 | + |
| 35 | +def resolve_optional(annotation: object) -> object: |
| 36 | + """Resolve an optional declaration to its non-None member declaration.""" |
| 37 | + |
| 38 | + member_type = optional_member_type(annotation) |
| 39 | + return annotation if member_type is None else member_type |
| 40 | + |
| 41 | + |
| 42 | +def resolve_annotated(annotation: object) -> object: |
| 43 | + """Resolve an ``Annotated[T, ...]`` declaration to its owned type ``T``.""" |
| 44 | + |
| 45 | + if get_origin(annotation) is Annotated: |
| 46 | + return get_args(annotation)[0] |
| 47 | + return annotation |
| 48 | + |
| 49 | + |
| 50 | +def is_enum_type(annotation: object) -> bool: |
| 51 | + """Return whether the declaration is an enum type.""" |
| 52 | + |
| 53 | + return isinstance(annotation, type) and issubclass(annotation, Enum) |
| 54 | + |
| 55 | + |
| 56 | +def enum_member_type(annotation: object) -> type[Enum] | None: |
| 57 | + """Return the enum type declared directly or as a simple optional.""" |
| 58 | + |
| 59 | + if is_enum_type(annotation): |
| 60 | + return annotation |
| 61 | + if not is_union_type(annotation): |
| 62 | + return None |
| 63 | + members = tuple( |
| 64 | + member for member in get_args(annotation) if member is not type(None) |
| 65 | + ) |
| 66 | + if len(members) != 1 or not is_enum_type(members[0]): |
| 67 | + return None |
| 68 | + return members[0] |
| 69 | + |
| 70 | + |
| 71 | +def is_list_of_enums(annotation: object) -> bool: |
| 72 | + """Return whether the declaration is ``list[SomeEnum]``.""" |
| 73 | + |
| 74 | + members = get_args(annotation) |
| 75 | + return ( |
| 76 | + get_origin(annotation) is list |
| 77 | + and len(members) == 1 |
| 78 | + and is_enum_type(members[0]) |
| 79 | + ) |
| 80 | + |
| 81 | + |
| 82 | +def get_enum_from_list(annotation: object) -> type[Enum] | None: |
| 83 | + """Return the enum type from ``list[SomeEnum]`` when declared.""" |
| 84 | + |
| 85 | + if not is_list_of_enums(annotation): |
| 86 | + return None |
| 87 | + return get_args(annotation)[0] |
0 commit comments