-
Notifications
You must be signed in to change notification settings - Fork 1
/
indent.py
88 lines (64 loc) · 2.19 KB
/
indent.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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import functools
import sys
from io import StringIO
import builtins
# this is the decorator that ensures all prints within the function indent by one
# extra level
def indenting(fn):
@functools.wraps(fn)
def wrapper(*args, **kwargs):
with indent:
return fn(*args, **kwargs)
return wrapper
_counter = 0 # ensure neighboring prints with same indent show divider
_level = 0 # start with no indent
# this class is the context manager for ensuring prints are indented
class indenter():
def __enter__(self):
global _level, _counter
_level += 1
_counter += 1
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
global _level, _counter
_level -= 1
_counter += 1
indent = indenter()
_no_end = False # to make sure we don't indent continued lines (prints following print(..., end='')
_last_info = (0, 0) # to check whether we have consecutive lines from two separate scopes at same indentation
def iprint(*args, indent=0, **kwargs):
global _no_end, _last_info
skip_tab, _no_end = _no_end, kwargs.get('end') == ''
if kwargs.get('file') or _level == 0 or skip_tab:
if _level == 0: _last_info = (_counter, _level)
return _print(*args, **kwargs)
io = StringIO()
_print(*args, file=io, **kwargs)
tabs = '\t' * (_level + indent)
info, _last_info = _last_info, (_counter, _level)
if info[0] != _last_info[0] and info[1] == _last_info[1]:
sys.stdout.write(tabs + '----\n')
out_str = ''.join(map(lambda l: tabs + l, io.getvalue().splitlines(True)))
sys.stdout.write(out_str)
_print = builtins.print # save old print function
builtins.print = iprint # override the default 'print' function
if __name__ == '__main__':
@indenting
def recur(i):
print(f'begin({i})')
if i > 0: recur(i - 1)
if i > 0: recur(i - 1)
print(f'end({i})')
with indent:
print('abc')
print('def')
with indent:
print('abc', end='')
print('def')
with indent:
print('abc')
print('def')
with indent:
print('abc')
print('def')
recur(2)