Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Utilities for union (sum type) disambiguation."""
3from __future__ import annotations
5from collections import defaultdict
6from collections.abc import Mapping
7from dataclasses import MISSING
8from functools import reduce
9from operator import or_
10from typing import TYPE_CHECKING, Any, Callable, Literal, Union, get_origin
12from attrs import NOTHING, Attribute, AttrsInstance
14from ._compat import (
15 NoneType,
16 adapted_fields,
17 fields_dict,
18 get_args,
19 has,
20 is_literal,
21 is_union_type,
22)
23from .gen import AttributeOverride
25if TYPE_CHECKING:
26 from .converters import BaseConverter
28__all__ = ["create_default_dis_func", "is_supported_union"]
31def is_supported_union(typ: Any) -> bool:
32 """Whether the type is a union of attrs classes or dataclasses."""
33 return is_union_type(typ) and all(
34 e is NoneType or has(get_origin(e) or e) for e in typ.__args__
35 )
38def create_default_dis_func(
39 converter: BaseConverter,
40 *classes: type[AttrsInstance],
41 use_literals: bool = True,
42 overrides: (
43 dict[str, AttributeOverride] | Literal["from_converter"]
44 ) = "from_converter",
45) -> Callable[[Mapping[Any, Any]], type[Any] | None]:
46 """Given attrs classes or dataclasses, generate a disambiguation function.
48 The function is based on unique fields without defaults or unique values.
50 :param use_literals: Whether to try using fields annotated as literals for
51 disambiguation.
52 :param overrides: Attribute overrides to apply.
54 .. versionchanged:: 24.1.0
55 Dataclasses are now supported.
56 """
57 if len(classes) < 2:
58 raise ValueError("At least two classes required.")
60 if overrides == "from_converter":
61 overrides = [
62 getattr(converter.get_structure_hook(c), "overrides", {}) for c in classes
63 ]
64 else:
65 overrides = [overrides for _ in classes]
67 # first, attempt for unique values
68 if use_literals:
69 # requirements for a discriminator field:
70 # (... TODO: a single fallback is OK)
71 # - it must always be enumerated
72 cls_candidates = [
73 {
74 at.name
75 for at in adapted_fields(get_origin(cl) or cl)
76 if is_literal(at.type)
77 }
78 for cl in classes
79 ]
81 # literal field names common to all members
82 discriminators: set[str] = cls_candidates[0]
83 for possible_discriminators in cls_candidates:
84 discriminators &= possible_discriminators
86 best_result = None
87 best_discriminator = None
88 for discriminator in discriminators:
89 # maps Literal values (strings, ints...) to classes
90 mapping = defaultdict(list)
92 for cl in classes:
93 for key in get_args(
94 fields_dict(get_origin(cl) or cl)[discriminator].type
95 ):
96 mapping[key].append(cl)
98 if best_result is None or max(len(v) for v in mapping.values()) <= max(
99 len(v) for v in best_result.values()
100 ):
101 best_result = mapping
102 best_discriminator = discriminator
104 if (
105 best_result
106 and best_discriminator
107 and max(len(v) for v in best_result.values()) != len(classes)
108 ):
109 final_mapping = {
110 k: v[0] if len(v) == 1 else Union[tuple(v)]
111 for k, v in best_result.items()
112 }
114 def dis_func(data: Mapping[Any, Any]) -> type | None:
115 if not isinstance(data, Mapping):
116 raise ValueError("Only input mappings are supported.")
117 return final_mapping[data[best_discriminator]]
119 return dis_func
121 # next, attempt for unique keys
123 # NOTE: This could just as well work with just field availability and not
124 # uniqueness, returning Unions ... it doesn't do that right now.
125 cls_and_attrs = [
126 (cl, *_usable_attribute_names(cl, override))
127 for cl, override in zip(classes, overrides)
128 ]
129 # For each class, attempt to generate a single unique required field.
130 uniq_attrs_dict: dict[str, type] = {}
132 # We start from classes with the largest number of unique fields
133 # so we can do easy picks first, making later picks easier.
134 remaining = sorted(cls_and_attrs, key=lambda c_a: len(c_a[1]), reverse=True)
136 # Iterate to a fixpoint. On each pass, pin every class that has a
137 # non-default field unique among the classes not yet pinned. Pinning a
138 # class removes its fields from consideration, which can make a previously
139 # ambiguous class uniquely identifiable, so we repeat until a pass pins
140 # nothing. A single pass in a fixed order would miss those later picks and
141 # fail to disambiguate valid unions depending on the class order.
142 made_progress = True
143 while made_progress:
144 made_progress = False
145 still_remaining = []
146 for entry in remaining:
147 cl, cl_reqs, back_map = entry
148 other_reqs = reduce(
149 or_, (c_a[1] for c_a in remaining if c_a[0] is not cl), set()
150 )
151 uniq = cl_reqs - other_reqs
153 # We want a unique attribute with no default.
154 cl_fields = fields_dict(get_origin(cl) or cl)
155 for maybe_renamed_attr_name in uniq:
156 orig_name = back_map[maybe_renamed_attr_name]
157 if cl_fields[orig_name].default in (NOTHING, MISSING):
158 uniq_attrs_dict[maybe_renamed_attr_name] = cl
159 made_progress = True
160 break
161 else:
162 still_remaining.append(entry)
163 remaining = still_remaining
165 # A class we could not pin to a unique non-default field can only be the
166 # single fallback (tried when no unique key matches). More than one such
167 # class is a genuine ambiguity.
168 if len(remaining) > 1:
169 raise TypeError(
170 "Could not disambiguate between "
171 + ", ".join(repr(c_a[0]) for c_a in remaining)
172 + ": no unique non-default attributes"
173 )
174 fallback = remaining[0][0] if remaining else None
176 if fallback is None:
178 def dis_func(data: Mapping[Any, Any]) -> type[AttrsInstance] | None:
179 if not isinstance(data, Mapping):
180 raise ValueError("Only input mappings are supported")
181 for k, v in uniq_attrs_dict.items():
182 if k in data:
183 return v
184 raise ValueError("Couldn't disambiguate")
186 else:
188 def dis_func(data: Mapping[Any, Any]) -> type[AttrsInstance] | None:
189 if not isinstance(data, Mapping):
190 raise ValueError("Only input mappings are supported")
191 for k, v in uniq_attrs_dict.items():
192 if k in data:
193 return v
194 return fallback
196 return dis_func
199create_uniq_field_dis_func = create_default_dis_func
202def _overriden_name(at: Attribute, override: AttributeOverride | None) -> str:
203 if override is None or override.rename is None:
204 return at.name
205 return override.rename
208def _usable_attribute_names(
209 cl: type[Any], overrides: dict[str, AttributeOverride]
210) -> tuple[set[str], dict[str, str]]:
211 """Return renamed fields and a mapping to original field names."""
212 res = set()
213 mapping = {}
215 for at in adapted_fields(get_origin(cl) or cl):
216 res.add(n := _overriden_name(at, overrides.get(at.name)))
217 mapping[n] = at.name
219 return res, mapping