Show More
@@ -100,23 +100,25 b' def extract_symbols(code, symbols):' | |||
|
100 | 100 | return [] |
|
101 | 101 | |
|
102 | 102 | marks = [(getattr(s, 'name', None), s.lineno) for s in py_code.body] |
|
103 | code = code.split('\n') | |
|
103 | 104 | |
|
104 | 105 | # construct a dictionary with elements |
|
105 | 106 | # {'symbol_name': (start_lineno, end_lineno), ...} |
|
106 |
end = |
|
|
107 | end = len(code) | |
|
107 | 108 | symbols_lines = {} |
|
108 | 109 | for name, start in reversed(marks): |
|
110 | while not code[end - 1].strip(): | |
|
111 | end -= 1 | |
|
109 | 112 | if name: |
|
110 | 113 | symbols_lines[name] = (start - 1, end) |
|
111 | 114 | end = start - 1 |
|
112 | 115 | |
|
113 | 116 | # fill a list with chunks of codes for each symbol |
|
114 | 117 | blocks = [] |
|
115 | code = code.split('\n') | |
|
116 | 118 | for symbol in symbols.split(','): |
|
117 | 119 | if symbol in symbols_lines: |
|
118 | 120 | start, end = symbols_lines[symbol] |
|
119 | blocks.append('\n'.join(code[start:end])) | |
|
121 | blocks.append('\n'.join(code[start:end]) + '\n') | |
|
120 | 122 | |
|
121 | 123 | return blocks |
|
122 | 124 |
@@ -65,9 +65,9 b' def test_extract_symbols():' | |||
|
65 | 65 | source = """import foo\na = 10\ndef b():\n return 42\nclass A: pass\n""" |
|
66 | 66 | symbols_args = ["a", "b", "A", "A,b", "A,a", "z"] |
|
67 | 67 | expected = [[], |
|
68 | ["def b():\n return 42"], | |
|
68 | ["def b():\n return 42\n"], | |
|
69 | 69 | ["class A: pass\n"], |
|
70 | ["class A: pass\n", "def b():\n return 42"], | |
|
70 | ["class A: pass\n", "def b():\n return 42\n"], | |
|
71 | 71 | ["class A: pass\n"], |
|
72 | 72 | []] |
|
73 | 73 | for symbols, exp in zip(symbols_args, expected): |
General Comments 0
You need to be logged in to leave comments.
Login now