diff --git a/IPython/core/magics/code.py b/IPython/core/magics/code.py index d23b4d3..ae9c45b 100644 --- a/IPython/core/magics/code.py +++ b/IPython/core/magics/code.py @@ -100,23 +100,25 @@ def extract_symbols(code, symbols): return [] marks = [(getattr(s, 'name', None), s.lineno) for s in py_code.body] + code = code.split('\n') # construct a dictionary with elements # {'symbol_name': (start_lineno, end_lineno), ...} - end = None + end = len(code) symbols_lines = {} for name, start in reversed(marks): + while not code[end - 1].strip(): + end -= 1 if name: symbols_lines[name] = (start - 1, end) end = start - 1 # fill a list with chunks of codes for each symbol blocks = [] - code = code.split('\n') for symbol in symbols.split(','): if symbol in symbols_lines: start, end = symbols_lines[symbol] - blocks.append('\n'.join(code[start:end])) + blocks.append('\n'.join(code[start:end]) + '\n') return blocks diff --git a/IPython/core/tests/test_magic.py b/IPython/core/tests/test_magic.py index f358882..e9bdd62 100644 --- a/IPython/core/tests/test_magic.py +++ b/IPython/core/tests/test_magic.py @@ -65,9 +65,9 @@ def test_extract_symbols(): source = """import foo\na = 10\ndef b():\n return 42\nclass A: pass\n""" symbols_args = ["a", "b", "A", "A,b", "A,a", "z"] expected = [[], - ["def b():\n return 42"], + ["def b():\n return 42\n"], ["class A: pass\n"], - ["class A: pass\n", "def b():\n return 42"], + ["class A: pass\n", "def b():\n return 42\n"], ["class A: pass\n"], []] for symbols, exp in zip(symbols_args, expected):