diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 62e3279eac1e..dab1fc23ef26 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -125,7 +125,9 @@ def transform(self) -> None: '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), + 'slots': _get_decorator_bool_argument(self._ctx, 'slots', False), } + py_version = self._ctx.api.options.python_version # 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 @@ -188,6 +190,9 @@ def transform(self) -> None: else: self._propertize_callables(attributes) + if decorator_arguments['slots']: + self.add_slots(info, attributes, correct_version=py_version >= (3, 10)) + self.reset_init_only_vars(info, attributes) self._add_dataclass_fields_magic_attribute() @@ -197,6 +202,35 @@ def transform(self) -> None: 'frozen': decorator_arguments['frozen'], } + def add_slots(self, + info: TypeInfo, + attributes: List[DataclassAttribute], + *, + correct_version: bool) -> None: + if not correct_version: + # This means that version is lower than `3.10`, + # it is just a non-existent argument for `dataclass` function. + self._ctx.api.fail( + 'Keyword argument "slots" for "dataclass" ' + 'is only valid in Python 3.10 and higher', + self._ctx.reason, + ) + return + if info.slots is not None or info.names.get('__slots__'): + # This means we have a slots confict. + # Class explicitly specifies `__slots__` field. + # And `@dataclass(slots=True)` is used. + # In runtime this raises a type error. + self._ctx.api.fail( + '"{}" both defines "__slots__" and is used with "slots=True"'.format( + self._ctx.cls.name, + ), + self._ctx.cls, + ) + return + + info.slots = {attr.name for attr in attributes} + def reset_init_only_vars(self, info: TypeInfo, attributes: List[DataclassAttribute]) -> None: """Remove init-only vars from the class and reset init var declarations.""" for attr in attributes: diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test index 2e5e1cfd5383..3d5f0a1da6e5 100644 --- a/test-data/unit/check-dataclasses.test +++ b/test-data/unit/check-dataclasses.test @@ -1350,3 +1350,74 @@ class Foo: reveal_type(Foo(bar=1.5)) # N: Revealed type is "__main__.Foo" [builtins fixtures/dataclasses.pyi] + + +[case testDataclassWithSlotsArg] +# flags: --python-version 3.10 +from dataclasses import dataclass + +@dataclass(slots=True) +class Some: + x: int + + def __init__(self, x: int) -> None: + self.x = x + self.y = 0 # E: Trying to assign name "y" that is not in "__slots__" of type "__main__.Some" + + def __post_init__(self) -> None: + self.y = 1 # E: Trying to assign name "y" that is not in "__slots__" of type "__main__.Some" +[builtins fixtures/dataclasses.pyi] + +[case testDataclassWithSlotsDef] +# flags: --python-version 3.10 +from dataclasses import dataclass + +@dataclass(slots=False) +class Some: + __slots__ = ('x',) + x: int + + def __init__(self, x: int) -> None: + self.x = x + self.y = 0 # E: Trying to assign name "y" that is not in "__slots__" of type "__main__.Some" + + def __post_init__(self) -> None: + self.y = 1 # E: Trying to assign name "y" that is not in "__slots__" of type "__main__.Some" +[builtins fixtures/dataclasses.pyi] + +[case testDataclassWithSlotsConflict] +# flags: --python-version 3.10 +from dataclasses import dataclass + +@dataclass(slots=True) +class Some: # E: "Some" both defines "__slots__" and is used with "slots=True" + __slots__ = ('x',) + x: int + +@dataclass(slots=True) +class EmptyDef: # E: "EmptyDef" both defines "__slots__" and is used with "slots=True" + __slots__ = () + x: int + +slots = ('x',) + +@dataclass(slots=True) +class DynamicDef: # E: "DynamicDef" both defines "__slots__" and is used with "slots=True" + __slots__ = slots + x: int +[builtins fixtures/dataclasses.pyi] + +[case testDataclassWithSlotsArgBefore310] +# flags: --python-version 3.9 +from dataclasses import dataclass + +@dataclass(slots=True) # E: Keyword argument "slots" for "dataclass" is only valid in Python 3.10 and higher +class Some: + x: int + +# Possible conflict: +@dataclass(slots=True) # E: Keyword argument "slots" for "dataclass" is only valid in Python 3.10 and higher +class Other: + __slots__ = ('x',) + x: int +[builtins fixtures/dataclasses.pyi]