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

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 try: 

42 # GenericMeta in py3.6 requires us to track __orig_bases__. This is fixed in py3.7 

43 # by the removal of GenericMeta. We should just be able to use cls.__bases__ in the 

44 # future. 

45 bases = getattr(cls, "__orig_bases__", cls.__bases__) 

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

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

48 cls = type(cls)(cls.__name__, bases, cls_dict) 

49 except TypeError: 

50 # We're in py3.7 and should use cls.__bases__ 

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

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

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

54 if qualname is not None: 

55 cls.__qualname__ = qualname 

56 

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

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

59 

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

61 return { 

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

63 for field in dataclasses.fields(self) 

64 if hasattr(self, field.name) 

65 } 

66 

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

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

69 object.__setattr__(self, fieldname, value) 

70 

71 cls.__getstate__ = __getstate__ 

72 cls.__setstate__ = __setstate__ 

73 

74 return cls