def test_generic_function_type(self): c = Callable([self.x, self.y], [ARG_POS, ARG_POS], [None, None], self.y, False, None, [TypeVarDef('X', -1, None)]) assert_equal(str(c), 'def [X] (X?, Y?) -> Y?') v = [TypeVarDef('Y', -1, None), TypeVarDef('X', -2, None)] c2 = Callable([], [], [], Void(None), False, None, v) assert_equal(str(c2), 'def [Y, X] ()')
def test_generic_function_type(self) -> None: c = CallableType([self.x, self.y], [ARG_POS, ARG_POS], [None, None], self.y, self.function, name=None, variables=[TypeVarDef('X', 'X', -1, [], self.fx.o)]) assert_equal(str(c), 'def [X] (X?, Y?) -> Y?') v = [TypeVarDef('Y', 'Y', -1, [], self.fx.o), TypeVarDef('X', 'X', -2, [], self.fx.o)] c2 = CallableType([], [], [], NoneType(), self.function, name=None, variables=v) assert_equal(str(c2), 'def [Y, X] ()')
def _add_cmp(ctx: 'mypy.plugin.ClassDefContext', adder: 'MethodAdder') -> None: """Generate all the cmp methods for this class.""" # For __ne__ and __eq__ the type is: # def __ne__(self, other: object) -> bool bool_type = ctx.api.named_type('__builtins__.bool') object_type = ctx.api.named_type('__builtins__.object') args = [Argument(Var('other', object_type), object_type, None, ARG_POS)] for method in ['__ne__', '__eq__']: adder.add_method(method, args, bool_type) # For the rest we use: # AT = TypeVar('AT') # def __lt__(self: AT, other: AT) -> bool # This way comparisons with subclasses will work correctly. tvd = TypeVarDef(SELF_TVAR_NAME, ctx.cls.info.fullname() + '.' + SELF_TVAR_NAME, -1, [], object_type) tvd_type = TypeVarType(tvd) self_tvar_expr = TypeVarExpr( SELF_TVAR_NAME, ctx.cls.info.fullname() + '.' + SELF_TVAR_NAME, [], object_type) ctx.cls.info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr) args = [Argument(Var('other', tvd_type), tvd_type, None, ARG_POS)] for method in ['__lt__', '__le__', '__gt__', '__ge__']: adder.add_method(method, args, bool_type, self_type=tvd_type, tvd=tvd)
def add_construct_method(self, fields: List["PydanticModelField"]) -> None: """ Adds a fully typed `construct` classmethod to the class. Similar to the fields-aware __init__ method, but always uses the field names (not aliases), and does not treat settings fields as optional. """ ctx = self._ctx set_str = ctx.api.named_type("__builtins__.set", [ctx.api.named_type("__builtins__.str")]) optional_set_str = UnionType([set_str, NoneType()]) fields_set_argument = Argument(Var("_fields_set", optional_set_str), optional_set_str, None, ARG_OPT) construct_arguments = self.get_field_arguments( fields, typed=True, force_all_optional=False, use_alias=False) construct_arguments = [fields_set_argument] + construct_arguments obj_type = ctx.api.named_type("__builtins__.object") self_tvar_name = "Model" tvar_fullname = ctx.cls.fullname + "." + self_tvar_name tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type) self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [], obj_type) ctx.cls.info.names[self_tvar_name] = SymbolTableNode( MDEF, self_tvar_expr) self_type = TypeVarType(tvd) add_method( ctx, "construct", construct_arguments, return_type=self_type, self_type=self_type, tvar_def=tvd, is_classmethod=True, )
def make_type_info(name: str, is_abstract: bool = False, mro: List[TypeInfo] = None, bases: List[Instance] = None, typevars: List[str] = None) -> TypeInfo: """Make a TypeInfo suitable for use in unit tests.""" type_def = TypeDef(name, Block([]), None, []) type_def.fullname = name if typevars: v = [] # type: List[TypeVarDef] id = 1 for n in typevars: v.append(TypeVarDef(n, id, None)) id += 1 type_def.type_vars = v info = TypeInfo(SymbolTable(), type_def) if mro is None: mro = [] info.mro = [info] + mro if bases is None: if mro: # By default, assume that there is a single non-generic base. bases = [Instance(mro[0], [])] else: bases = [] info.bases = bases return info
def add_class_tvars(t: Type, itype: Instance, is_classmethod: bool, builtin_type: Callable[[str], Instance], original_type: Type) -> Type: """Instantiate type variables during analyze_class_attribute_access, e.g T and Q in the following: def A(Generic(T)): @classmethod def foo(cls: Type[Q]) -> Tuple[T, Q]: ... class B(A): pass B.foo() original_type is the value of the type B in the expression B.foo() """ # TODO: verify consistency between Q and T info = itype.type # type: TypeInfo if isinstance(t, CallableType): # TODO: Should we propagate type variable values? tvars = [TypeVarDef(n, n, i + 1, [], builtin_type('builtins.object'), tv.variance) for (i, n), tv in zip(enumerate(info.type_vars), info.defn.type_vars)] if is_classmethod: t = bind_self(t, original_type, is_classmethod=True) return t.copy_modified(variables=tvars + t.variables) elif isinstance(t, Overloaded): return Overloaded([cast(CallableType, add_class_tvars(item, itype, is_classmethod, builtin_type, original_type)) for item in t.items()]) return t
def anal_var_defs(self, var_defs: List[TypeVarDef]) -> List[TypeVarDef]: a = [] # type: List[TypeVarDef] for vd in var_defs: a.append( TypeVarDef(vd.name, vd.id.raw_id, self.anal_array(vd.values), vd.upper_bound.accept(self), vd.variance, vd.line)) return a
def add_class_tvars(t: Type, info: TypeInfo, is_classmethod: bool, builtin_type: Callable[[str], Instance]) -> Type: if isinstance(t, CallableType): # TODO: Should we propagate type variable values? vars = [ TypeVarDef(n, i + 1, None, builtin_type('builtins.object'), tv.variance) for (i, n), tv in zip(enumerate(info.type_vars), info.defn.type_vars) ] arg_types = t.arg_types arg_kinds = t.arg_kinds arg_names = t.arg_names if is_classmethod: arg_types = arg_types[1:] arg_kinds = arg_kinds[1:] arg_names = arg_names[1:] return t.copy_modified(arg_types=arg_types, arg_kinds=arg_kinds, arg_names=arg_names, variables=vars + t.variables) elif isinstance(t, Overloaded): return Overloaded([ cast(CallableType, add_class_tvars(i, info, is_classmethod, builtin_type)) for i in t.items() ]) return t
def anal_var_defs(self, var_defs: List[TypeVarDef]) -> List[TypeVarDef]: a = List[TypeVarDef]() for vd in var_defs: a.append( TypeVarDef(vd.name, vd.id, self.anal_array(vd.values), vd.upper_bound.accept(self), vd.line, vd.repr)) return a
def add_class_tvars(t: Type, info: TypeInfo, is_classmethod: bool, builtin_type: Callable[[str], Instance]) -> Type: if isinstance(t, CallableType): # TODO: Should we propagate type variable values? vars = [TypeVarDef(n, i + 1, None, builtin_type('builtins.object')) for i, n in enumerate(info.type_vars)] arg_types = t.arg_types arg_kinds = t.arg_kinds arg_names = t.arg_names if is_classmethod: arg_types = arg_types[1:] arg_kinds = arg_kinds[1:] arg_names = arg_names[1:] return CallableType(arg_types, arg_kinds, arg_names, t.ret_type, t.fallback, t.name, vars + t.variables, t.bound_vars, t.line) elif isinstance(t, Overloaded): return Overloaded([cast(CallableType, add_class_tvars(i, info, is_classmethod, builtin_type)) for i in t.items()]) return t
def bind_new(self, name: str, tvar_expr: TypeVarLikeExpr) -> TypeVarLikeDef: if self.is_class_scope: self.class_id += 1 i = self.class_id else: self.func_id -= 1 i = self.func_id if isinstance(tvar_expr, TypeVarExpr): tvar_def = TypeVarDef( name, tvar_expr.fullname, i, values=tvar_expr.values, upper_bound=tvar_expr.upper_bound, variance=tvar_expr.variance, line=tvar_expr.line, column=tvar_expr.column ) # type: TypeVarLikeDef elif isinstance(tvar_expr, ParamSpecExpr): tvar_def = ParamSpecDef( name, tvar_expr.fullname, i, line=tvar_expr.line, column=tvar_expr.column ) else: assert False self.scope[tvar_expr.fullname] = tvar_def return tvar_def
def add_class_tvars(t: Type, itype: Instance, isuper: Optional[Instance], is_classmethod: bool, builtin_type: Callable[[str], Instance], original_type: Type) -> Type: """Instantiate type variables during analyze_class_attribute_access, e.g T and Q in the following: class A(Generic[T]): @classmethod def foo(cls: Type[Q]) -> Tuple[T, Q]: ... class B(A[str]): pass B.foo() original_type is the value of the type B in the expression B.foo() """ # TODO: verify consistency between Q and T info = itype.type # type: TypeInfo if is_classmethod: assert isuper is not None t = expand_type_by_instance(t, isuper) # We add class type variables if the class method is accessed on class object # without applied type arguments, this matches the behavior of __init__(). # For example (continuing the example in docstring): # A # The type of callable is def [T] () -> A[T], _not_ def () -> A[Any] # A[int] # The type of callable is def () -> A[int] # and # A.foo # The type is generic def [T] () -> Tuple[T, A[T]] # A[int].foo # The type is non-generic def () -> Tuple[int, A[int]] # # This behaviour is useful for defining alternative constructors for generic classes. # To achieve such behaviour, we add the class type variables that are still free # (i.e. appear in the return type of the class object on which the method was accessed). free_ids = {t.id for t in itype.args if isinstance(t, TypeVarType)} if isinstance(t, CallableType): # NOTE: in practice either all or none of the variables are free, since # visit_type_application() will detect any type argument count mismatch and apply # a correct number of Anys. tvars = [ TypeVarDef(n, n, i + 1, [], builtin_type('builtins.object'), tv.variance) for (i, n), tv in zip(enumerate(info.type_vars), info.defn.type_vars) # use 'is' to avoid id clashes with unrelated variables if any(tv.id is id for id in free_ids) ] if is_classmethod: t = bind_self(t, original_type, is_classmethod=True) return t.copy_modified(variables=tvars + t.variables) elif isinstance(t, Overloaded): return Overloaded([ cast( CallableType, add_class_tvars(item, itype, isuper, is_classmethod, builtin_type, original_type)) for item in t.items() ]) return t
def _type_var_def(name: str, module: str, upper_bound, values=(), meta_level=0) -> TypeVarDef: id_ = TypeVarId.new(meta_level) id_.raw_id = -id_.raw_id fullname = f'{module}.{name}' return TypeVarDef(name, fullname, id_, list(values), upper_bound)
def translate_variables(self, variables: List[TypeVarDef]) -> List[TypeVarDef]: if not variables: return variables items = [] # type: List[TypeVarDef] for v in variables: if v.id > 0: items.append( TypeVarDef(v.name, -v.id - self.num_func_tvars, v.values)) else: items.append(v) return items
def callable(self, vars, *a) -> Callable: """callable(args, a1, ..., an, r) constructs a callable with argument types a1, ... an and return type r and type arguments vars. """ tv = [] # type: List[TypeVarDef] n = -1 for v in vars: tv.append(TypeVarDef(v, n, None)) n -= 1 return Callable(a[:-1], [ARG_POS] * (len(a) - 1), [None] * (len(a) - 1), a[-1], False, None, tv)
def add_class_tvars(t: Type, info: TypeInfo) -> Type: if isinstance(t, Callable): vars = [ TypeVarDef(n, i + 1, None) for i, n in enumerate(info.type_vars) ] return Callable(t.arg_types, t.arg_kinds, t.arg_names, t.ret_type, t.is_type_obj(), t.name, vars + t.variables, t.bound_vars, t.line, None) elif isinstance(t, Overloaded): return Overloaded( [cast(Callable, add_class_tvars(i, info)) for i in t.items()]) return t
def class_callable(init_type: Callable, info: TypeInfo) -> Callable: """Create a type object type based on the signature of __init__.""" variables = [] # type: List[TypeVarDef] for i, tvar in enumerate(info.defn.type_vars): variables.append(TypeVarDef(tvar.name, i + 1, tvar.values)) initvars = init_type.variables variables.extend(initvars) c = Callable(init_type.arg_types, init_type.arg_kinds, init_type.arg_names, self_type(info), True, None, variables).with_name('"{}"'.format(info.name())) return convert_class_tvars_to_func_tvars(c, len(initvars))
def bind(self, name: str, tvar_expr: TypeVarExpr) -> TypeVarDef: if self.is_class_scope: self.class_id += 1 i = self.class_id else: self.func_id -= 1 i = self.func_id tvar_def = TypeVarDef( name, i, values=tvar_expr.values, upper_bound=tvar_expr.upper_bound, variance=tvar_expr.variance, line=tvar_expr.line, column=tvar_expr.column) self.scope[tvar_expr.fullname()] = tvar_def return tvar_def
def class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance) -> CallableType: """Create a type object type based on the signature of __init__.""" variables = [] # type: List[TypeVarDef] for i, tvar in enumerate(info.defn.type_vars): variables.append(TypeVarDef(tvar.name, i + 1, tvar.values, tvar.upper_bound, tvar.variance)) initvars = init_type.variables variables.extend(initvars) callable_type = init_type.copy_modified( ret_type=self_type(info), fallback=type_type, name=None, variables=variables) c = callable_type.with_name('"{}"'.format(info.name())) return convert_class_tvars_to_func_tvars(c, len(initvars))
def callable(self, vars: List[str], *a: Type) -> CallableType: """callable(args, a1, ..., an, r) constructs a callable with argument types a1, ... an and return type r and type arguments vars. """ tv = [] # type: List[TypeVarDef] n = -1 for v in vars: tv.append(TypeVarDef(v, v, n, [], self.fx.o)) n -= 1 return CallableType(list(a[:-1]), [ARG_POS] * (len(a) - 1), [None] * (len(a) - 1), a[-1], self.fx.function, name=None, variables=tv)
def freshen_function_type_vars(callee: F) -> F: """Substitute fresh type variables for generic function type variables.""" if isinstance(callee, CallableType): if not callee.is_generic(): return cast(F, callee) tvdefs = [] tvmap = {} # type: Dict[TypeVarId, Type] for v in callee.variables: tvdef = TypeVarDef.new_unification_variable(v) tvdefs.append(tvdef) tvmap[v.id] = TypeVarType(tvdef) fresh = cast(CallableType, expand_type(callee, tvmap)).copy_modified(variables=tvdefs) return cast(F, fresh) else: assert isinstance(callee, Overloaded) fresh_overload = Overloaded([freshen_function_type_vars(item) for item in callee.items()]) return cast(F, fresh_overload)
def _add_cmp(ctx: 'mypy.plugin.ClassDefContext', adder: 'MethodAdder') -> None: """Generate all the cmp methods for this class.""" # For __ne__ and __eq__ the type is: # def __ne__(self, other: object) -> bool bool_type = ctx.api.named_type('__builtins__.bool') object_type = ctx.api.named_type('__builtins__.object') args = [Argument(Var('other', object_type), object_type, None, ARG_POS)] for method in ['__ne__', '__eq__']: adder.add_method(method, args, bool_type) # For the rest we use: # AT = TypeVar('AT') # def __lt__(self: AT, other: AT) -> bool # This way comparisons with subclasses will work correctly. tvd = TypeVarDef('AT', 'AT', 1, [], object_type) tvd_type = TypeVarType(tvd) args = [Argument(Var('other', tvd_type), tvd_type, None, ARG_POS)] for method in ['__lt__', '__le__', '__gt__', '__ge__']: adder.add_method(method, args, bool_type, self_type=tvd_type, tvd=tvd)
def make_type_info(self, name: str, module_name: str = None, is_abstract: bool = False, mro: List[TypeInfo] = None, bases: List[Instance] = None, typevars: List[str] = None, variances: List[int] = None) -> TypeInfo: """Make a TypeInfo suitable for use in unit tests.""" class_def = ClassDef(name, Block([]), None, []) class_def.fullname = name if module_name is None: if '.' in name: module_name = name.rsplit('.', 1)[0] else: module_name = '__main__' if typevars: v = [] # type: List[TypeVarDef] for id, n in enumerate(typevars, 1): if variances: variance = variances[id - 1] else: variance = COVARIANT v.append(TypeVarDef(n, id, None, self.o, variance=variance)) class_def.type_vars = v info = TypeInfo(SymbolTable(), class_def, module_name) if mro is None: mro = [] if name != 'builtins.object': mro.append(self.oi) info.mro = [info] + mro if bases is None: if mro: # By default, assume that there is a single non-generic base. bases = [Instance(mro[0], [])] else: bases = [] info.bases = bases return info
def add_construct_method(self, fields: List['PydanticModelField']) -> None: """ Adds a fully typed `construct` classmethod to the class. Similar to the fields-aware __init__ method, but always uses the field names (not aliases), and does not treat settings fields as optional. """ ctx = self._ctx set_str = ctx.api.named_type( f'{BUILTINS_NAME}.set', [ctx.api.named_type(f'{BUILTINS_NAME}.str')]) optional_set_str = UnionType([set_str, NoneType()]) fields_set_argument = Argument(Var('_fields_set', optional_set_str), optional_set_str, None, ARG_OPT) construct_arguments = self.get_field_arguments( fields, typed=True, force_all_optional=False, use_alias=False) construct_arguments = [fields_set_argument] + construct_arguments obj_type = ctx.api.named_type(f'{BUILTINS_NAME}.object') self_tvar_name = '_PydanticBaseModel' # Make sure it does not conflict with other names in the class tvar_fullname = ctx.cls.fullname + '.' + self_tvar_name tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type) self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [], obj_type) ctx.cls.info.names[self_tvar_name] = SymbolTableNode( MDEF, self_tvar_expr) # Backward-compatible with TypeVarDef from Mypy 0.910. if isinstance(tvd, TypeVarType): self_type = tvd else: self_type = TypeVarType(tvd) # type: ignore[call-arg] add_method( ctx, 'construct', construct_arguments, return_type=self_type, self_type=self_type, tvar_def=tvd, is_classmethod=True, )
def test_type_variable_binding(self) -> None: assert_equal(str(TypeVarDef('X', 'X', 1, [], self.fx.o)), 'X') assert_equal(str(TypeVarDef('X', 'X', 1, [self.x, self.y], self.fx.o)), 'X in (X?, Y?)')
def transform(self) -> None: """Apply all the necessary transformations to the underlying dataclass so as to ensure it is fully type checked according to the rules in PEP 557. """ ctx = self._ctx info = self._ctx.cls.info attributes = self.collect_attributes() if attributes is None: # Some definitions are not ready, defer() should be already called. return for attr in attributes: if attr.type is None: ctx.api.defer() return decorator_arguments = { 'init': _get_decorator_bool_argument(self._ctx, 'init', True), 'eq': _get_decorator_bool_argument(self._ctx, 'eq', True), 'order': _get_decorator_bool_argument(self._ctx, 'order', False), 'frozen': _get_decorator_bool_argument(self._ctx, 'frozen', False), } if info.get('replace') is None: obj_type = ctx.api.named_type('__builtins__.object') self_tvar_expr = TypeVarExpr(SELF_UVAR_NAME, info.fullname + '.' + SELF_UVAR_NAME, [], obj_type) info.names[SELF_UVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr) replace_tvar_def = TypeVarDef(SELF_UVAR_NAME, info.fullname + '.' + SELF_UVAR_NAME, -1, [], fill_typevars(info)) replace_other_type = TypeVarType(replace_tvar_def) add_method(ctx, 'replace', args=[ Argument( Var('changes', AnyType(TypeOfAny.explicit)), AnyType(TypeOfAny.explicit), None, ARG_STAR2) ], return_type=replace_other_type, self_type=replace_other_type, tvar_def=replace_tvar_def) # If there are no attributes, it may be that the semantic analyzer has not # processed them yet. In order to work around this, we can simply skip generating # __init__ if there are no attributes, because if the user truly did not define any, # then the object default __init__ with an empty signature will be present anyway. if (decorator_arguments['init'] and ('__init__' not in info.names or info.names['__init__'].plugin_generated) and attributes): add_method( ctx, '__init__', args=[ attr.to_argument() for attr in attributes if attr.is_in_init ], return_type=NoneType(), ) if (decorator_arguments['eq'] and info.get('__eq__') is None or decorator_arguments['order']): # Type variable for self types in generated methods. obj_type = ctx.api.named_type('__builtins__.object') self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, [], obj_type) info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr) # Add <, >, <=, >=, but only if the class has an eq method. if decorator_arguments['order']: if not decorator_arguments['eq']: ctx.api.fail('eq must be True if order is True', ctx.cls) for method_name in ['__lt__', '__gt__', '__le__', '__ge__']: # Like for __eq__ and __ne__, we want "other" to match # the self type. obj_type = ctx.api.named_type('__builtins__.object') order_tvar_def = TypeVarDef( SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, -1, [], obj_type) order_other_type = TypeVarType(order_tvar_def) order_return_type = ctx.api.named_type('__builtins__.bool') order_args = [ Argument(Var('other', order_other_type), order_other_type, None, ARG_POS) ] existing_method = info.get(method_name) if existing_method is not None and not existing_method.plugin_generated: assert existing_method.node ctx.api.fail( 'You may not have a custom %s method when order=True' % method_name, existing_method.node, ) add_method( ctx, method_name, args=order_args, return_type=order_return_type, self_type=order_other_type, tvar_def=order_tvar_def, ) if decorator_arguments['frozen']: self._freeze(attributes) self.reset_init_only_vars(info, attributes) info.metadata['dataclass'] = { 'attributes': [attr.serialize() for attr in attributes], 'frozen': decorator_arguments['frozen'], }
def transform(self) -> None: """Apply all the necessary transformations to the underlying dataclass so as to ensure it is fully type checked according to the rules in PEP 557. """ ctx = self._ctx info = self._ctx.cls.info attributes = self.collect_attributes() decorator_arguments = { 'init': _get_decorator_bool_argument(self._ctx, 'init', True), 'eq': _get_decorator_bool_argument(self._ctx, 'eq', True), 'order': _get_decorator_bool_argument(self._ctx, 'order', False), 'frozen': _get_decorator_bool_argument(self._ctx, 'frozen', False), } if decorator_arguments['init']: _add_method( ctx, '__init__', args=[ attr.to_argument(info) for attr in attributes if attr.is_in_init ], return_type=NoneTyp(), ) for stmt in self._ctx.cls.defs.body: # Fix up the types of classmethods since, by default, # they will be based on the parent class' init. if isinstance(stmt, Decorator) and stmt.func.is_class: func_type = stmt.func.type if isinstance(func_type, CallableType): func_type.arg_types[0] = self._ctx.api.class_type( self._ctx.cls.info) if isinstance(stmt, OverloadedFuncDef) and stmt.is_class: func_type = stmt.type if isinstance(func_type, Overloaded): class_type = ctx.api.class_type(ctx.cls.info) for item in func_type.items(): item.arg_types[0] = class_type if stmt.impl is not None: assert isinstance(stmt.impl, Decorator) if isinstance(stmt.impl.func.type, CallableType): stmt.impl.func.type.arg_types[0] = class_type # Add an eq method, but only if the class doesn't already have one. if decorator_arguments['eq'] and info.get('__eq__') is None: for method_name in ['__eq__', '__ne__']: # The TVar is used to enforce that "other" must have # the same type as self (covariant). Note the # "self_type" parameter to _add_method. obj_type = ctx.api.named_type('__builtins__.object') cmp_tvar_def = TypeVarDef('T', 'T', -1, [], obj_type) cmp_other_type = TypeVarType(cmp_tvar_def) cmp_return_type = ctx.api.named_type('__builtins__.bool') _add_method( ctx, method_name, args=[ Argument(Var('other', cmp_other_type), cmp_other_type, None, ARG_POS) ], return_type=cmp_return_type, self_type=cmp_other_type, tvar_def=cmp_tvar_def, ) # Add <, >, <=, >=, but only if the class has an eq method. if decorator_arguments['order']: if not decorator_arguments['eq']: ctx.api.fail('eq must be True if order is True', ctx.cls) for method_name in ['__lt__', '__gt__', '__le__', '__ge__']: # Like for __eq__ and __ne__, we want "other" to match # the self type. obj_type = ctx.api.named_type('__builtins__.object') order_tvar_def = TypeVarDef('T', 'T', -1, [], obj_type) order_other_type = TypeVarType(order_tvar_def) order_return_type = ctx.api.named_type('__builtins__.bool') order_args = [ Argument(Var('other', order_other_type), order_other_type, None, ARG_POS) ] existing_method = info.get(method_name) if existing_method is not None: assert existing_method.node ctx.api.fail( 'You may not have a custom %s method when order=True' % method_name, existing_method.node, ) _add_method( ctx, method_name, args=order_args, return_type=order_return_type, self_type=order_other_type, tvar_def=order_tvar_def, ) if decorator_arguments['frozen']: self._freeze(attributes) # Remove init-only vars from the class. for attr in attributes: if attr.is_init_var: del info.names[attr.name] info.metadata['dataclass'] = { 'attributes': OrderedDict((attr.name, attr.serialize()) for attr in attributes), 'frozen': decorator_arguments['frozen'], }
def build_namedtuple_typeinfo(self, name: str, items: List[str], types: List[Type], default_items: Mapping[str, Expression], line: int) -> TypeInfo: strtype = self.api.named_type('__builtins__.str') implicit_any = AnyType(TypeOfAny.special_form) basetuple_type = self.api.named_type('__builtins__.tuple', [implicit_any]) dictype = (self.api.named_type_or_none('builtins.dict', [strtype, implicit_any]) or self.api.named_type('__builtins__.object')) # Actual signature should return OrderedDict[str, Union[types]] ordereddictype = (self.api.named_type_or_none('builtins.dict', [strtype, implicit_any]) or self.api.named_type('__builtins__.object')) fallback = self.api.named_type('__builtins__.tuple', [implicit_any]) # Note: actual signature should accept an invariant version of Iterable[UnionType[types]]. # but it can't be expressed. 'new' and 'len' should be callable types. iterable_type = self.api.named_type_or_none('typing.Iterable', [implicit_any]) function_type = self.api.named_type('__builtins__.function') info = self.api.basic_new_typeinfo(name, fallback) info.is_named_tuple = True tuple_base = TupleType(types, fallback) info.tuple_type = tuple_base info.line = line # We can't calculate the complete fallback type until after semantic # analysis, since otherwise base classes might be incomplete. Postpone a # callback function that patches the fallback. self.api.schedule_patch(PRIORITY_FALLBACKS, lambda: calculate_tuple_fallback(tuple_base)) def add_field(var: Var, is_initialized_in_class: bool = False, is_property: bool = False) -> None: var.info = info var.is_initialized_in_class = is_initialized_in_class var.is_property = is_property var._fullname = '%s.%s' % (info.fullname, var.name) info.names[var.name] = SymbolTableNode(MDEF, var) fields = [Var(item, typ) for item, typ in zip(items, types)] for var in fields: add_field(var, is_property=True) # We can't share Vars between fields and method arguments, since they # have different full names (the latter are normally used as local variables # in functions, so their full names are set to short names when generated methods # are analyzed). vars = [Var(item, typ) for item, typ in zip(items, types)] tuple_of_strings = TupleType([strtype for _ in items], basetuple_type) add_field(Var('_fields', tuple_of_strings), is_initialized_in_class=True) add_field(Var('_field_types', dictype), is_initialized_in_class=True) add_field(Var('_field_defaults', dictype), is_initialized_in_class=True) add_field(Var('_source', strtype), is_initialized_in_class=True) add_field(Var('__annotations__', ordereddictype), is_initialized_in_class=True) add_field(Var('__doc__', strtype), is_initialized_in_class=True) tvd = TypeVarDef(SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, -1, [], info.tuple_type) selftype = TypeVarType(tvd) def add_method(funcname: str, ret: Type, args: List[Argument], is_classmethod: bool = False, is_new: bool = False, ) -> None: if is_classmethod or is_new: first = [Argument(Var('_cls'), TypeType.make_normalized(selftype), None, ARG_POS)] else: first = [Argument(Var('_self'), selftype, None, ARG_POS)] args = first + args types = [arg.type_annotation for arg in args] items = [arg.variable.name for arg in args] arg_kinds = [arg.kind for arg in args] assert None not in types signature = CallableType(cast(List[Type], types), arg_kinds, items, ret, function_type) signature.variables = [tvd] func = FuncDef(funcname, args, Block([])) func.info = info func.is_class = is_classmethod func.type = set_callable_name(signature, func) func._fullname = info.fullname + '.' + funcname func.line = line if is_classmethod: v = Var(funcname, func.type) v.is_classmethod = True v.info = info v._fullname = func._fullname func.is_decorated = True dec = Decorator(func, [NameExpr('classmethod')], v) dec.line = line sym = SymbolTableNode(MDEF, dec) else: sym = SymbolTableNode(MDEF, func) sym.plugin_generated = True info.names[funcname] = sym add_method('_replace', ret=selftype, args=[Argument(var, var.type, EllipsisExpr(), ARG_NAMED_OPT) for var in vars]) def make_init_arg(var: Var) -> Argument: default = default_items.get(var.name, None) kind = ARG_POS if default is None else ARG_OPT return Argument(var, var.type, default, kind) add_method('__new__', ret=selftype, args=[make_init_arg(var) for var in vars], is_new=True) add_method('_asdict', args=[], ret=ordereddictype) special_form_any = AnyType(TypeOfAny.special_form) add_method('_make', ret=selftype, is_classmethod=True, args=[Argument(Var('iterable', iterable_type), iterable_type, None, ARG_POS), Argument(Var('new'), special_form_any, EllipsisExpr(), ARG_NAMED_OPT), Argument(Var('len'), special_form_any, EllipsisExpr(), ARG_NAMED_OPT)]) self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, [], info.tuple_type) info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr) return info
def make_type_var(name: str, id: int, values: List[Type], upper_bound: Type, variance: int) -> TypeVarType: return TypeVarType( TypeVarDef(name, id, values, upper_bound, variance))
def transform(self) -> None: """Apply all the necessary transformations to the underlying dataclass so as to ensure it is fully type checked according to the rules in PEP 557. """ ctx = self._ctx info = self._ctx.cls.info attributes = self.collect_attributes() if attributes is None: # Some definitions are not ready, defer() should be already called. return for attr in attributes: if attr.type is None: ctx.api.defer() return decorator_arguments = { "init": _get_decorator_bool_argument(self._ctx, "init", True), "eq": _get_decorator_bool_argument(self._ctx, "eq", True), "order": _get_decorator_bool_argument(self._ctx, "order", False), "frozen": _get_decorator_bool_argument(self._ctx, "frozen", False), } # If there are no attributes, it may be that the semantic analyzer has not # processed them yet. In order to work around this, we can simply skip # generating __init__ if there are no attributes, because if the user # truly did not define any, then the object default __init__ with an # empty signature will be present anyway. if (decorator_arguments["init"] and ("__init__" not in info.names or info.names["__init__"].plugin_generated) and attributes): add_method( ctx, "__init__", args=[ attr.to_argument() for attr in attributes if attr.is_in_init ], return_type=NoneType(), ) if (decorator_arguments["eq"] and info.get("__eq__") is None or decorator_arguments["order"]): # Type variable for self types in generated methods. obj_type = ctx.api.named_type("__builtins__.object") self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname + "." + SELF_TVAR_NAME, [], obj_type) info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr) # Add <, >, <=, >=, but only if the class has an eq method. if decorator_arguments["order"]: if not decorator_arguments["eq"]: ctx.api.fail("eq must be True if order is True", ctx.cls) for method_name in ["__lt__", "__gt__", "__le__", "__ge__"]: # Like for __eq__ and __ne__, we want "other" to match # the self type. obj_type = ctx.api.named_type("__builtins__.object") order_tvar_def = TypeVarDef( SELF_TVAR_NAME, info.fullname + "." + SELF_TVAR_NAME, -1, [], obj_type, ) order_other_type = TypeVarType(order_tvar_def) order_return_type = ctx.api.named_type("__builtins__.bool") order_args = [ Argument(Var("other", order_other_type), order_other_type, None, ARG_POS) ] existing_method = info.get(method_name) if existing_method is not None and not existing_method.plugin_generated: assert existing_method.node ctx.api.fail( "You may not have a custom %s method when order=True" % method_name, existing_method.node, ) add_method( ctx, method_name, args=order_args, return_type=order_return_type, self_type=order_other_type, tvar_def=order_tvar_def, ) if decorator_arguments["frozen"]: self._freeze(attributes) self.reset_init_only_vars(info, attributes) info.metadata["dataclass"] = { "attributes": [attr.serialize() for attr in attributes], "frozen": decorator_arguments["frozen"], }