|
1 | 1 | """ |
2 | 2 | Module used for managing annotations. |
3 | 3 | """ |
4 | | -from typing import Union, Optional, get_args, Generic, get_origin, get_type_hints |
| 4 | +from typing import Union, Optional, Generic, Iterable, get_args, get_origin, get_type_hints |
5 | 5 | from datetime import datetime, timedelta, timezone |
| 6 | +from inspect import isclass, isabstract |
| 7 | +from itertools import product, chain |
6 | 8 | from contextlib import suppress |
7 | | -from inspect import isclass |
8 | | -from .doc import doc_category |
| 9 | + |
9 | 10 | from .utilities import issubclass_noexcept |
| 11 | +from .doc import doc_category |
10 | 12 |
|
11 | 13 |
|
12 | 14 | __all__ = ( |
13 | 15 | "register_annotations", |
14 | 16 | "get_annotations", |
| 17 | + "convert_types", |
15 | 18 | ) |
16 | 19 |
|
17 | 20 |
|
@@ -116,3 +119,74 @@ def get_annotations(class_) -> dict: |
116 | 119 | del annotations["return"] |
117 | 120 |
|
118 | 121 | return annotations |
| 122 | + |
| 123 | + |
| 124 | +def convert_types(input_type: type): |
| 125 | + """ |
| 126 | + Type preprocessing method, that extends the list of types with inherited members (polymorphism) |
| 127 | + and removes classes that are wrapped by some other class, if the wrapper class also appears in |
| 128 | + the annotations. |
| 129 | + """ |
| 130 | + def remove_classes(types: list): |
| 131 | + r = types.copy() |
| 132 | + for type_ in types: |
| 133 | + # It's a wrapper of some class -> remove the wrapped class |
| 134 | + if hasattr(type_, "__wrapped__"): |
| 135 | + if type_.__wrapped__ in r: |
| 136 | + r.remove(type_.__wrapped__) |
| 137 | + |
| 138 | + # Abstract classes are classes that don't allow instantiation -> remove the class |
| 139 | + if isabstract(type_): |
| 140 | + r.remove(type_) |
| 141 | + |
| 142 | + return tuple({a:0 for a in r}) |
| 143 | + |
| 144 | + |
| 145 | + if isinstance(input_type, str): |
| 146 | + raise TypeError( |
| 147 | + f"Provided type '{input_type}' is not a type - it is a string!\n" |
| 148 | + "Potential subscripted type problem?\n" |
| 149 | + "Instead of e. g., list['type'], try using typing.List['type']." |
| 150 | + ) |
| 151 | + |
| 152 | + origin = get_origin(input_type) |
| 153 | + if issubclass_noexcept(origin, Generic): |
| 154 | + # Patch for Python versions < 3.10 |
| 155 | + input_type.__name__ = origin.__name__ |
| 156 | + |
| 157 | + # Unpack Union items into a tuple |
| 158 | + if origin is Union or issubclass_noexcept(origin, (Iterable, Generic)): |
| 159 | + new_types = [] |
| 160 | + for arg_group in get_args(input_type): |
| 161 | + new_types.append(remove_classes(list(convert_types(arg_group)))) |
| 162 | + |
| 163 | + if origin is Union: |
| 164 | + return tuple(chain.from_iterable(new_types)) # Just expand unions |
| 165 | + |
| 166 | + # Process abstract classes and polymorphism |
| 167 | + new_origins = [] |
| 168 | + for origin in convert_types(origin): |
| 169 | + if issubclass_noexcept(origin, Generic): |
| 170 | + for comb in product(*new_types): |
| 171 | + new_origins.append(origin[comb]) |
| 172 | + elif issubclass_noexcept(origin, Iterable): |
| 173 | + new = origin[tuple(chain.from_iterable(new_types))] if len(new_types) else origin |
| 174 | + new_origins.append(new) |
| 175 | + else: |
| 176 | + new_origins.append(origin) |
| 177 | + |
| 178 | + return remove_classes(new_origins) |
| 179 | + |
| 180 | + if input_type.__module__ == "builtins": |
| 181 | + # Don't consider built-int types for polymorphism |
| 182 | + # No removal of abstract classes is needed either as builtins types aren't abstract |
| 183 | + return (input_type,) |
| 184 | + |
| 185 | + # Extend subclasses |
| 186 | + subtypes = [] |
| 187 | + if hasattr(input_type, "__subclasses__"): |
| 188 | + for st in input_type.__subclasses__(): |
| 189 | + subtypes.extend(convert_types(st)) |
| 190 | + |
| 191 | + # Remove wrapped classes (eg. wrapped by decorator) + ABC classes |
| 192 | + return remove_classes([input_type, *subtypes]) |
0 commit comments