Skip to content

Commit e717eb1

Browse files
committed
feat: deprecate union aliasing on Python 3.14 and later
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
1 parent 6640341 commit e717eb1

File tree

1 file changed

+172
-126
lines changed

1 file changed

+172
-126
lines changed

plum/alias.py

Lines changed: 172 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -26,139 +26,185 @@
2626
parsing how unions print.
2727
"""
2828

29+
import sys
2930
from functools import wraps
3031
from typing import TypeVar, Union, _type_repr, get_args
3132

33+
from typing_extensions import deprecated
34+
3235
__all__ = ["activate_union_aliases", "deactivate_union_aliases", "set_union_alias"]
3336

3437
UnionT = TypeVar("UnionT")
3538

36-
_union_type = type(Union[int, float])
37-
_original_repr = _union_type.__repr__
38-
_original_str = _union_type.__str__
39-
40-
41-
@wraps(_original_repr)
42-
def _new_repr(self: object) -> str:
43-
"""Print a `typing.Union`, replacing all aliased unions by their aliased names.
44-
45-
Returns:
46-
str: Representation of a `typing.Union` taking into account union aliases.
47-
"""
48-
args = get_args(self)
49-
args_set = set(args)
50-
51-
# Find all aliased unions contained in this union.
52-
found_unions = []
53-
found_positions = []
54-
found_aliases = []
55-
for union, alias in reversed(_ALIASED_UNIONS):
56-
union_set = set(union)
57-
if union_set <= args_set:
58-
found = False
59-
for i, arg in enumerate(args):
60-
if arg in union_set:
61-
found_unions.append(union_set)
62-
found_positions.append(i)
63-
found_aliases.append(alias)
64-
found = True
39+
40+
if sys.version_info < (3, 14):
41+
_union_type = type(Union[int, float])
42+
_original_repr = _union_type.__repr__
43+
_original_str = _union_type.__str__
44+
45+
@wraps(_original_repr)
46+
def _new_repr(self: object) -> str:
47+
"""Print a `typing.Union`, replacing all aliased unions by their aliased names.
48+
49+
Returns:
50+
str: Representation of a `typing.Union` taking into account union aliases.
51+
"""
52+
args = get_args(self)
53+
args_set = set(args)
54+
55+
# Find all aliased unions contained in this union.
56+
found_unions = []
57+
found_positions = []
58+
found_aliases = []
59+
for union, alias in reversed(_ALIASED_UNIONS):
60+
union_set = set(union)
61+
if union_set <= args_set:
62+
found = False
63+
for i, arg in enumerate(args):
64+
if arg in union_set:
65+
found_unions.append(union_set)
66+
found_positions.append(i)
67+
found_aliases.append(alias)
68+
found = True
69+
break
70+
if not found: # pragma: no cover
71+
# This branch should never be reached.
72+
raise AssertionError(
73+
"Could not identify union. This should never happen."
74+
)
75+
76+
# Delete any unions that are contained in strictly bigger unions. We check for
77+
# strictly inequality because any union includes itself.
78+
for i in range(len(found_unions) - 1, -1, -1):
79+
for union in found_unions:
80+
if found_unions[i] < union:
81+
del found_unions[i]
82+
del found_positions[i]
83+
del found_aliases[i]
6584
break
66-
if not found: # pragma: no cover
67-
# This branch should never be reached.
68-
raise AssertionError(
69-
"Could not identify union. This should never happen."
70-
)
7185

72-
# Delete any unions that are contained in strictly bigger unions. We check for
73-
# strictly inequality because any union includes itself.
74-
for i in range(len(found_unions) - 1, -1, -1):
86+
# Create a set with all arguments of all found unions.
87+
found_args = set()
7588
for union in found_unions:
76-
if found_unions[i] < union:
77-
del found_unions[i]
78-
del found_positions[i]
79-
del found_aliases[i]
80-
break
81-
82-
# Create a set with all arguments of all found unions.
83-
found_args = set()
84-
for union in found_unions:
85-
found_args |= union
86-
87-
# Insert the aliases right before the first found argument. When we insert an
88-
# element, the positions of following insertions need to be appropriately
89-
# incremented.
90-
args = list(args)
91-
# Sort by insertion position to ensure that all following insertions are at higher
92-
# indices. This makes the bookkeeping simple.
93-
for delta, (i, alias) in enumerate(
94-
sorted(zip(found_positions, found_aliases), key=lambda x: x[0])
95-
):
96-
args.insert(i + delta, alias)
97-
98-
# Filter all elements of unions that are aliased.
99-
new_args = ()
100-
for arg in args:
101-
if arg not in found_args:
102-
new_args += (arg,)
103-
args = new_args
104-
105-
# Generate a string representation.
106-
args_repr = [a if isinstance(a, str) else _type_repr(a) for a in args]
107-
# Like `typing` does, print `Optional` whenever possible.
108-
if len(args) == 2:
109-
if args[0] is type(None): # noqa: E721
110-
return f"typing.Optional[{args_repr[1]}]"
111-
elif args[1] is type(None): # noqa: E721
112-
return f"typing.Optional[{args_repr[0]}]"
113-
# We would like to just print `args_repr[0]` whenever `len(args) == 1`, but
114-
# this might break code that parses how unions print.
115-
return "typing.Union[" + ", ".join(args_repr) + "]"
116-
117-
118-
@wraps(_original_str)
119-
def _new_str(self: object) -> str:
120-
"""Does the same as :func:`_new_repr`.
121-
122-
Returns:
123-
str: Representation of the `typing.Union` taking into account union aliases.
124-
"""
125-
return _new_repr(self)
126-
127-
128-
def activate_union_aliases() -> None:
129-
"""When printing `typing.Union`s, replace all aliased unions by the aliased names.
130-
This monkey patches `__repr__` and `__str__` for `typing.Union`."""
131-
_union_type.__repr__ = _new_repr
132-
_union_type.__str__ = _new_str
133-
134-
135-
def deactivate_union_aliases() -> None:
136-
"""Undo what :func:`.alias.activate` did. This restores the original `__repr__`
137-
and `__str__` for `typing.Union`."""
138-
_union_type.__repr__ = _original_repr
139-
_union_type.__str__ = _original_str
140-
141-
142-
_ALIASED_UNIONS: list = []
143-
144-
145-
def set_union_alias(union: UnionT, alias: str) -> UnionT:
146-
"""Change how a `typing.Union` is printed. This does not modify `union`.
147-
148-
Args:
149-
union (type or type hint): A union.
150-
alias (str): How to print `union`.
151-
152-
Returns:
153-
type or type hint: `union`.
154-
"""
155-
args = get_args(union) if isinstance(union, _union_type) else (union,)
156-
for existing_union, existing_alias in _ALIASED_UNIONS:
157-
if set(existing_union) == set(args) and alias != existing_alias:
158-
if isinstance(union, _union_type):
159-
union_str = _original_str(union)
160-
else:
161-
union_str = repr(union)
162-
raise RuntimeError(f"`{union_str}` already has alias `{existing_alias}`.")
163-
_ALIASED_UNIONS.append((args, alias))
164-
return union
89+
found_args |= union
90+
91+
# Insert the aliases right before the first found argument. When we insert an
92+
# element, the positions of following insertions need to be appropriately
93+
# incremented.
94+
args = list(args)
95+
# Sort by insertion position to ensure that all following insertions are
96+
# at higher indices. This makes the bookkeeping simple.
97+
for delta, (i, alias) in enumerate(
98+
sorted(zip(found_positions, found_aliases), key=lambda x: x[0])
99+
):
100+
args.insert(i + delta, alias)
101+
102+
# Filter all elements of unions that are aliased.
103+
new_args = ()
104+
for arg in args:
105+
if arg not in found_args:
106+
new_args += (arg,)
107+
args = new_args
108+
109+
# Generate a string representation.
110+
args_repr = [a if isinstance(a, str) else _type_repr(a) for a in args]
111+
# Like `typing` does, print `Optional` whenever possible.
112+
if len(args) == 2:
113+
if args[0] is type(None): # noqa: E721
114+
return f"typing.Optional[{args_repr[1]}]"
115+
elif args[1] is type(None): # noqa: E721
116+
return f"typing.Optional[{args_repr[0]}]"
117+
# We would like to just print `args_repr[0]` whenever `len(args) == 1`, but
118+
# this might break code that parses how unions print.
119+
return "typing.Union[" + ", ".join(args_repr) + "]"
120+
121+
@wraps(_original_str)
122+
def _new_str(self: object) -> str:
123+
"""Does the same as :func:`_new_repr`.
124+
125+
Returns:
126+
str: Representation of the `typing.Union` taking into account union aliases.
127+
"""
128+
return _new_repr(self)
129+
130+
def activate_union_aliases() -> None:
131+
"""When printing `typing.Union`s, replace aliased unions by the aliased names.
132+
This monkey patches `__repr__` and `__str__` for `typing.Union`."""
133+
_union_type.__repr__ = _new_repr
134+
_union_type.__str__ = _new_str
135+
136+
def deactivate_union_aliases() -> None:
137+
"""Undo what :func:`.alias.activate` did. This restores the original `__repr__`
138+
and `__str__` for `typing.Union`."""
139+
_union_type.__repr__ = _original_repr
140+
_union_type.__str__ = _original_str
141+
142+
_ALIASED_UNIONS: list = []
143+
144+
def set_union_alias(union: UnionT, alias: str) -> UnionT:
145+
"""Change how a `typing.Union` is printed. This does not modify `union`.
146+
147+
Args:
148+
union (type or type hint): A union.
149+
alias (str): How to print `union`.
150+
151+
Returns:
152+
type or type hint: `union`.
153+
"""
154+
if sys.version_info >= (3, 14):
155+
return union
156+
157+
args = get_args(union) if isinstance(union, _union_type) else (union,)
158+
for existing_union, existing_alias in _ALIASED_UNIONS:
159+
if set(existing_union) == set(args) and alias != existing_alias:
160+
if isinstance(union, _union_type):
161+
union_str = _original_str(union)
162+
else:
163+
union_str = repr(union)
164+
raise RuntimeError(
165+
f"`{union_str}` already has alias `{existing_alias}`."
166+
)
167+
_ALIASED_UNIONS.append((args, alias))
168+
return union
169+
170+
171+
else:
172+
173+
@deprecated(
174+
"Plum's union aliasing is not supported on Python 3.14 and later.",
175+
category=RuntimeWarning,
176+
stacklevel=2,
177+
)
178+
def activate_union_aliases() -> None:
179+
"""When printing `typing.Union`s, replace aliased unions by the aliased names.
180+
This monkey patches `__repr__` and `__str__` for `typing.Union`."""
181+
pass
182+
183+
@deprecated(
184+
"Plum's union aliasing is not supported on Python 3.14 and later.",
185+
category=RuntimeWarning,
186+
stacklevel=2,
187+
)
188+
def deactivate_union_aliases() -> None:
189+
"""Undo what :func:`.alias.activate` did. This restores the original `__repr__`
190+
and `__str__` for `typing.Union`."""
191+
if sys.version_info < (3, 14):
192+
_union_type.__repr__ = _original_repr
193+
_union_type.__str__ = _original_str
194+
195+
@deprecated(
196+
"Plum's union aliasing is not supported on Python 3.14 and later.",
197+
category=RuntimeWarning,
198+
stacklevel=2,
199+
)
200+
def set_union_alias(union: UnionT, alias: str) -> UnionT:
201+
"""Change how a `typing.Union` is printed. This does not modify `union`.
202+
203+
Args:
204+
union (type or type hint): A union.
205+
alias (str): How to print `union`.
206+
207+
Returns:
208+
type or type hint: `union`.
209+
"""
210+
return union

0 commit comments

Comments
 (0)