71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
|
||
|
|
# Author: Leonardo Gama (@leogama)
|
||
|
|
# Copyright (c) 2022-2023 The Uncertainty Quantification Foundation.
|
||
|
|
# License: 3-clause BSD. The full license text is available at:
|
||
|
|
# - https://github.com/uqfoundation/dill/blob/master/LICENSE
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import re
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
import dill
|
||
|
|
from dill import detect
|
||
|
|
from dill.logger import stderr_handler, adapter as logger
|
||
|
|
|
||
|
|
try:
|
||
|
|
from StringIO import StringIO
|
||
|
|
except ImportError:
|
||
|
|
from io import StringIO
|
||
|
|
|
||
|
|
test_obj = {'a': (1, 2), 'b': object(), 'f': lambda x: x**2, 'big': list(range(10))}
|
||
|
|
|
||
|
|
def test_logging(should_trace):
|
||
|
|
buffer = StringIO()
|
||
|
|
handler = logging.StreamHandler(buffer)
|
||
|
|
logger.addHandler(handler)
|
||
|
|
try:
|
||
|
|
dill.dumps(test_obj)
|
||
|
|
if should_trace:
|
||
|
|
regex = re.compile(r'(\S*┬ \w.*[^)]' # begin pickling object
|
||
|
|
r'|│*└ # \w.* \[\d+ (\wi)?B])' # object written (with size)
|
||
|
|
)
|
||
|
|
for line in buffer.getvalue().splitlines():
|
||
|
|
assert regex.fullmatch(line)
|
||
|
|
return buffer.getvalue()
|
||
|
|
else:
|
||
|
|
assert buffer.getvalue() == ""
|
||
|
|
finally:
|
||
|
|
logger.removeHandler(handler)
|
||
|
|
buffer.close()
|
||
|
|
|
||
|
|
def test_trace_to_file(stream_trace):
|
||
|
|
file = tempfile.NamedTemporaryFile(mode='r')
|
||
|
|
with detect.trace(file.name, mode='w'):
|
||
|
|
dill.dumps(test_obj)
|
||
|
|
file_trace = file.read()
|
||
|
|
file.close()
|
||
|
|
# Apparently, objects can change location in memory...
|
||
|
|
reghex = re.compile(r'0x[0-9A-Za-z]+')
|
||
|
|
file_trace, stream_trace = reghex.sub('0x', file_trace), reghex.sub('0x', stream_trace)
|
||
|
|
# PyPy prints dictionary contents with repr(dict)...
|
||
|
|
regdict = re.compile(r'(dict\.__repr__ of ).*')
|
||
|
|
file_trace, stream_trace = regdict.sub(r'\1{}>', file_trace), regdict.sub(r'\1{}>', stream_trace)
|
||
|
|
assert file_trace == stream_trace
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
logger.removeHandler(stderr_handler)
|
||
|
|
test_logging(should_trace=False)
|
||
|
|
detect.trace(True)
|
||
|
|
test_logging(should_trace=True)
|
||
|
|
detect.trace(False)
|
||
|
|
test_logging(should_trace=False)
|
||
|
|
|
||
|
|
loglevel = logging.ERROR
|
||
|
|
logger.setLevel(loglevel)
|
||
|
|
with detect.trace():
|
||
|
|
stream_trace = test_logging(should_trace=True)
|
||
|
|
test_logging(should_trace=False)
|
||
|
|
assert logger.getEffectiveLevel() == loglevel
|
||
|
|
test_trace_to_file(stream_trace)
|