Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/libcst/_add_slots.py: 85%

1# This file is derived from github.com/ericvsmith/dataclasses, and is Apache 2 licensed. 

2# https://github.com/ericvsmith/dataclasses/blob/ae712dd993420d43444f188f452/LICENSE.txt 

3# https://github.com/ericvsmith/dataclasses/blob/ae712dd993420d43444f/dataclass_tools.py 

4# Changed: takes slots in base classes into account when creating slots 

5 

6import dataclasses 

7from itertools import chain, filterfalse 

8from typing import Any, Mapping, Type, TypeVar 

9 

10_T = TypeVar("_T") 

11 

12 

13def add_slots(cls: Type[_T]) -> Type[_T]: 

14 # Need to create a new class, since we can't set __slots__ 

15 # after a class has been created. 

16 

17 # Make sure __slots__ isn't already set. 

18 if "__slots__" in cls.__dict__: 

19 raise TypeError(f"{cls.__name__} already specifies __slots__") 

20 

21 # Create a new dict for our new class. 

22 cls_dict = dict(cls.__dict__) 

23 field_names = tuple(f.name for f in dataclasses.fields(cls)) 

24 inherited_slots = set( 

25 chain.from_iterable( 

26 superclass.__dict__.get("__slots__", ()) for superclass in cls.mro() 

27 ) 

28 ) 

29 cls_dict["__slots__"] = tuple( 

30 filterfalse(inherited_slots.__contains__, field_names) 

31 ) 

32 for field_name in field_names: 

33 # Remove our attributes, if present. They'll still be 

34 # available in _MARKER. 

35 cls_dict.pop(field_name, None) 

36 # Remove __dict__ itself. 

37 cls_dict.pop("__dict__", None) 

38 

39 # Create the class. 

40 qualname = getattr(cls, "__qualname__", None) 

41 

42 # pyre-fixme[9]: cls has type `Type[Variable[_T]]`; used as `_T`. 

43 # pyre-fixme[19]: Expected 0 positional arguments. 

44 cls = type(cls)(cls.__name__, cls.__bases__, cls_dict) 

45 if qualname is not None: 

46 cls.__qualname__ = qualname 

47 

48 # Set __getstate__ and __setstate__ to workaround a bug with pickling frozen 

49 # dataclasses with slots. See https://bugs.python.org/issue36424 

50 

51 def __getstate__(self: object) -> Mapping[str, Any]: 

52 return { 

53 field.name: getattr(self, field.name) 

54 for field in dataclasses.fields(self) 

55 if hasattr(self, field.name) 

56 } 

57 

58 def __setstate__(self: object, state: Mapping[str, Any]) -> None: 

59 for fieldname, value in state.items(): 

60 object.__setattr__(self, fieldname, value) 

61 

62 cls.__getstate__ = __getstate__ 

63 cls.__setstate__ = __setstate__ 

64 

65 return cls