forked from bloomberg/memray
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
60 lines (49 loc) · 1.84 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
"""Utilities / Helpers for writing tests."""
from dataclasses import dataclass
from typing import List
from typing import Optional
from typing import Tuple
from memray import AllocatorType
def filter_relevant_allocations(records, ranged=False):
addresses = set()
filter_allocations = [AllocatorType.VALLOC]
filter_deallocations = [AllocatorType.FREE]
if ranged:
filter_allocations.append(AllocatorType.MMAP)
filter_deallocations.append(AllocatorType.MUNMAP)
for record in records:
if record.allocator in filter_allocations:
yield record
addresses.add(record.address)
elif record.allocator in filter_deallocations:
if record.address in addresses:
yield record
addresses.discard(record.address)
@dataclass
class MockAllocationRecord:
"""Mimics :py:class:`memray._memray.AllocationRecord`."""
tid: int
address: int
size: int
allocator: AllocatorType
stack_id: int
n_allocations: int
_stack: Optional[List[Tuple[str, str, int]]] = None
_hybrid_stack: Optional[List[Tuple[str, str, int]]] = None
@property
def thread_name(self):
return str(hex(self.tid)) if self.tid != -1 else "merged thread"
@staticmethod
def __get_stack_trace(stack, max_stacks):
if max_stacks == 0:
return stack
else:
return stack[:max_stacks]
def stack_trace(self, max_stacks=0):
if self._stack is None:
raise AssertionError("did not expect a call to `stack_trace`")
return self.__get_stack_trace(self._stack, max_stacks)
def hybrid_stack_trace(self, max_stacks=0):
if self._hybrid_stack is None:
raise AssertionError("did not expect a call to `hybrid_stack_trace`")
return self.__get_stack_trace(self._hybrid_stack, max_stacks)