Spaces:
Sleeping
Sleeping
| import re | |
| from typing import Any, Optional, Type, TypeVar | |
| from dataclasses import is_dataclass | |
| import attr | |
| from attr import NOTHING, resolve_types | |
| from ._compat import get_args, get_origin, is_generic, adapted_fields | |
| class AttributeOverride(object): | |
| omit_if_default: Optional[bool] = attr.ib(default=None) | |
| rename: Optional[str] = attr.ib(default=None) | |
| def override(omit_if_default=None, rename=None): | |
| return AttributeOverride(omit_if_default=omit_if_default, rename=rename) | |
| _neutral = AttributeOverride() | |
| def make_dict_unstructure_fn(cl, converter, omit_if_default=False, **kwargs): | |
| """Generate a specialized dict unstructuring function for an attrs class.""" | |
| cl_name = cl.__name__ | |
| fn_name = "unstructure_" + cl_name | |
| globs = {} | |
| lines = [] | |
| post_lines = [] | |
| attrs = adapted_fields(cl) # type: ignore | |
| lines.append(f"def {fn_name}(i):") | |
| lines.append(" res = {") | |
| for a in attrs: | |
| attr_name = a.name | |
| override = kwargs.pop(attr_name, _neutral) | |
| kn = attr_name if override.rename is None else override.rename | |
| d = a.default | |
| # For each attribute, we try resolving the type here and now. | |
| # If a type is manually overwritten, this function should be | |
| # regenerated. | |
| if a.type is not None: | |
| handler = converter._unstructure_func.dispatch(a.type) | |
| else: | |
| handler = converter.unstructure | |
| is_identity = handler == converter._unstructure_identity | |
| if not is_identity: | |
| unstruct_handler_name = f"__cattr_unstruct_handler_{attr_name}" | |
| globs[unstruct_handler_name] = handler | |
| invoke = f"{unstruct_handler_name}(i.{attr_name})" | |
| else: | |
| invoke = f"i.{attr_name}" | |
| if d is not attr.NOTHING and ( | |
| (omit_if_default and override.omit_if_default is not False) | |
| or override.omit_if_default | |
| ): | |
| def_name = f"__cattr_def_{attr_name}" | |
| if isinstance(d, attr.Factory): | |
| globs[def_name] = d.factory | |
| if d.takes_self: | |
| post_lines.append( | |
| f" if i.{attr_name} != {def_name}(i):" | |
| ) | |
| else: | |
| post_lines.append(f" if i.{attr_name} != {def_name}():") | |
| post_lines.append(f" res['{kn}'] = {invoke}") | |
| else: | |
| globs[def_name] = d | |
| post_lines.append(f" if i.{attr_name} != {def_name}:") | |
| post_lines.append(f" res['{kn}'] = {invoke}") | |
| else: | |
| # No default or no override. | |
| lines.append(f" '{kn}': {invoke},") | |
| lines.append(" }") | |
| total_lines = lines + post_lines + [" return res"] | |
| eval(compile("\n".join(total_lines), "", "exec"), globs) | |
| fn = globs[fn_name] | |
| return fn | |
| def generate_mapping(cl: Type, old_mapping): | |
| mapping = {} | |
| for p, t in zip(get_origin(cl).__parameters__, get_args(cl)): | |
| if isinstance(t, TypeVar): | |
| continue | |
| mapping[p.__name__] = t | |
| if not mapping: | |
| return old_mapping | |
| cls = attr.make_class( | |
| "GenericMapping", | |
| {x: attr.attrib() for x in mapping.keys()}, | |
| frozen=True, | |
| ) | |
| return cls(**mapping) | |
| def make_dict_structure_fn( | |
| cl: Type, converter, _cattrs_forbid_extra_keys: bool = False, **kwargs | |
| ): | |
| """Generate a specialized dict structuring function for an attrs class.""" | |
| mapping = None | |
| if is_generic(cl): | |
| base = get_origin(cl) | |
| mapping = generate_mapping(cl, mapping) | |
| cl = base | |
| for base in getattr(cl, "__orig_bases__", ()): | |
| if is_generic(base) and not str(base).startswith("typing.Generic"): | |
| mapping = generate_mapping(base, mapping) | |
| break | |
| if isinstance(cl, TypeVar): | |
| cl = getattr(mapping, cl.__name__, cl) | |
| cl_name = cl.__name__ | |
| fn_name = "structure_" + cl_name | |
| # We have generic parameters and need to generate a unique name for the function | |
| for p in getattr(cl, "__parameters__", ()): | |
| # This is nasty, I am not sure how best to handle `typing.List[str]` or `TClass[int, int]` as a parameter type here | |
| name_base = getattr(mapping, p.__name__) | |
| name = getattr(name_base, "__name__", str(name_base)) | |
| name = re.sub(r"[\[\.\] ,]", "_", name) | |
| fn_name += f"_{name}" | |
| globs = {"__c_s": converter.structure, "__cl": cl, "__m": mapping} | |
| lines = [] | |
| post_lines = [] | |
| attrs = adapted_fields(cl) | |
| is_dc = is_dataclass(cl) | |
| if any(isinstance(a.type, str) for a in attrs): | |
| # PEP 563 annotations - need to be resolved. | |
| resolve_types(cl) | |
| lines.append(f"def {fn_name}(o, *_):") | |
| lines.append(" res = {") | |
| for a in attrs: | |
| an = a.name | |
| override = kwargs.pop(an, _neutral) | |
| type = a.type | |
| if isinstance(type, TypeVar): | |
| type = getattr(mapping, type.__name__, type) | |
| # For each attribute, we try resolving the type here and now. | |
| # If a type is manually overwritten, this function should be | |
| # regenerated. | |
| if type is not None: | |
| handler = converter._structure_func.dispatch(type) | |
| else: | |
| handler = converter.structure | |
| struct_handler_name = f"__cattr_struct_handler_{an}" | |
| globs[struct_handler_name] = handler | |
| ian = an if (is_dc or an[0] != "_") else an[1:] | |
| kn = an if override.rename is None else override.rename | |
| globs[f"__c_t_{an}"] = type | |
| if a.default is NOTHING: | |
| lines.append( | |
| f" '{ian}': {struct_handler_name}(o['{kn}'], __c_t_{an})," | |
| ) | |
| else: | |
| post_lines.append(f" if '{kn}' in o:") | |
| post_lines.append( | |
| f" res['{ian}'] = {struct_handler_name}(o['{kn}'], __c_t_{an})" | |
| ) | |
| lines.append(" }") | |
| if _cattrs_forbid_extra_keys: | |
| allowed_fields = {a.name for a in attrs} | |
| globs["__c_a"] = allowed_fields | |
| post_lines += [ | |
| " unknown_fields = set(o.keys()) - __c_a", | |
| " if unknown_fields:", | |
| " raise Exception(", | |
| f" 'Extra fields in constructor for {cl_name}: ' + ', '.join(unknown_fields)" | |
| " )", | |
| ] | |
| total_lines = lines + post_lines + [" return __cl(**res)"] | |
| eval(compile("\n".join(total_lines), "", "exec"), globs) | |
| return globs[fn_name] | |
| def make_iterable_unstructure_fn(cl: Any, converter, unstructure_to=None): | |
| """Generate a specialized unstructure function for an iterable.""" | |
| handler = converter.unstructure | |
| fn_name = "unstructure_iterable" | |
| # Let's try fishing out the type args. | |
| if getattr(cl, "__args__", None) is not None: | |
| type_arg = get_args(cl)[0] | |
| # We can do the dispatch here and now. | |
| handler = converter._unstructure_func.dispatch(type_arg) | |
| globs = {"__cattr_seq_cl": unstructure_to or cl, "__cattr_u": handler} | |
| lines = [] | |
| lines.append(f"def {fn_name}(iterable):") | |
| lines.append(" res = __cattr_seq_cl(__cattr_u(i) for i in iterable)") | |
| total_lines = lines + [" return res"] | |
| eval(compile("\n".join(total_lines), "", "exec"), globs) | |
| fn = globs[fn_name] | |
| return fn | |
| def make_mapping_unstructure_fn(cl: Any, converter, unstructure_to=None): | |
| """Generate a specialized unstructure function for a mapping.""" | |
| key_handler = converter.unstructure | |
| val_handler = converter.unstructure | |
| fn_name = "unstructure_mapping" | |
| # Let's try fishing out the type args. | |
| if getattr(cl, "__args__", None) is not None: | |
| args = get_args(cl) | |
| if len(args) == 2: | |
| key_arg, val_arg = args | |
| else: | |
| # Probably a Counter | |
| key_arg, val_arg = args, Any | |
| # We can do the dispatch here and now. | |
| key_handler = converter._unstructure_func.dispatch(key_arg) | |
| if key_handler == converter._unstructure_identity: | |
| key_handler = None | |
| val_handler = converter._unstructure_func.dispatch(val_arg) | |
| if val_handler == converter._unstructure_identity: | |
| val_handler = None | |
| globs = { | |
| "__cattr_mapping_cl": unstructure_to or cl, | |
| "__cattr_k_u": key_handler, | |
| "__cattr_v_u": val_handler, | |
| } | |
| if key_handler is not None: | |
| globs["__cattr_k_u"] | |
| if val_handler is not None: | |
| globs["__cattr_v_u"] | |
| k_u = "__cattr_k_u(k)" if key_handler is not None else "k" | |
| v_u = "__cattr_v_u(v)" if val_handler is not None else "v" | |
| lines = [] | |
| lines.append(f"def {fn_name}(mapping):") | |
| lines.append( | |
| f" res = __cattr_mapping_cl(({k_u}, {v_u}) for k, v in mapping.items())" | |
| ) | |
| total_lines = lines + [" return res"] | |
| eval(compile("\n".join(total_lines), "", "exec"), globs) | |
| fn = globs[fn_name] | |
| return fn | |