Coverage for common/memsize.py: 91%

1""" 

2crate_anon/common/memsize.py 

3 

4=============================================================================== 

5 

6 Copyright (C) 2015, University of Cambridge, Department of Psychiatry. 

7 Created by Rudolf Cardinal (rnc1001@cam.ac.uk). 

8 

9 This file is part of CRATE. 

10 

11 CRATE is free software: you can redistribute it and/or modify 

12 it under the terms of the GNU General Public License as published by 

13 the Free Software Foundation, either version 3 of the License, or 

14 (at your option) any later version. 

15 

16 CRATE is distributed in the hope that it will be useful, 

17 but WITHOUT ANY WARRANTY; without even the implied warranty of 

18 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

19 GNU General Public License for more details. 

20 

21 You should have received a copy of the GNU General Public License 

22 along with CRATE. If not, see <https://www.gnu.org/licenses/>. 

23 

24=============================================================================== 

25 

26Calculate the size of objects in memory (fast). 

27 

28From https://stackoverflow.com/questions/449560/how-do-i-determine-the-size-of-an-object-in-python 

29 

30""" # noqa: E501 

31 

32from gc import get_referents 

33from sys import getsizeof 

34from types import ModuleType, FunctionType 

35from typing import Any, List, Set 

36 

37# Custom objects know their class. 

38# Function objects seem to know way too much, including modules. 

39# Exclude modules as well. 

40DENYLIST = type, ModuleType, FunctionType 

41 

42 

43def getsize(obj: Any, assume_none_denied: bool = False) -> int: 

44 """ 

45 Return the total size (in bytes) of the object and its members. 

46 From https://stackoverflow.com/questions/449560/how-do-i-determine-the-size-of-an-object-in-python 

47 

48 Args: 

49 obj: 

50 Object to check. 

51 assume_none_denied: 

52 Skip checks for classes/modules/functions. Assume that all objects 

53 should be checked (typically, meaning that the caller guarantees 

54 not to pass stuff that doesn't need checking). 

55 """ # noqa: E501 

56 if not assume_none_denied: 

57 if isinstance(obj, DENYLIST): 

58 raise TypeError( 

59 f"getsize() does not take argument of type: " f"{type(obj)}" 

60 ) 

61 seen_ids = set() # type: Set[int] 

62 size = 0 

63 objects = [obj] # type: List[Any] 

64 while objects: 

65 need_referents = [] # type: List[Any] 

66 for obj in objects: 

67 if assume_none_denied or not isinstance(obj, DENYLIST): 

68 obj_id = id(obj) 

69 if obj_id not in seen_ids: 

70 seen_ids.add(obj_id) 

71 size += getsizeof(obj) 

72 need_referents.append(obj) 

73 objects = get_referents(*need_referents) 

74 return size