%load -s to load specific functions or classes

IPython/core/magics/code.py

@@ -18,6 +18,7 @@
 import os
 import re
 import sys
+import ast
 from itertools import chain
 
 # Our own packages
@@ -30,6 +31,7 @@
 from IPython.utils.contexts import preserve_keys
 from IPython.utils.path import get_py_filename, unquote_filename
 from IPython.utils.warn import warn
+from IPython.utils.text import get_text_list
 
 #-----------------------------------------------------------------------------
 # Magic implementation classes
@@ -77,6 +79,58 @@ def extract_code_ranges(ranges_str):
         yield (start, end)
 
 
+@skip_doctest
+def extract_symbols(code, symbols):
+    """
+    Return a tuple  (blocks, not_found)
+    where ``blocks`` is a list of code fragments
+    for each symbol parsed from code, and ``not_found`` are
+    symbols not found in the code.
+
+    For example::
+
+        >>> code = '''a = 10
+
+        def b(): return 42
+
+        class A: pass'''
+
+        >>> extract_symbols(code, 'A,b')
+        (["class A: pass", "def b(): return 42"], [])
+    """
+    try:
+        py_code = ast.parse(code)
+    except SyntaxError:
+        # ignores non python code
+        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 = 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 = []
+    not_found = []
+    for symbol in symbols.split(','):
+        if symbol in symbols_lines:
+            start, end = symbols_lines[symbol]
+            blocks.append('\n'.join(code[start:end]) + '\n')
+        else:
+            not_found.append(symbol)
+
+    return blocks, not_found
+
+
 class InteractivelyDefined(Exception):
     """Exception for interactively defined variable in magic_edit"""
     def __init__(self, index):
@@ -217,6 +271,8 @@ def load(self, arg_s):
           (x..(y-1)). Both limits x and y can be left blank (meaning the 
           beginning and end of the file, respectively).
 
+          -s <symbols>: Specify function or classes to load from python source. 
+
           -y : Don't ask confirmation for loading source above 200 000 characters.
 
         This magic command can either take a local filename, a URL, an history
@@ -230,15 +286,27 @@ def load(self, arg_s):
         %load http://www.example.com/myscript.py
         %load -r 5-10 myscript.py
         %load -r 10-20,30,40: foo.py
+        %load -s MyClass,wonder_function myscript.py
         """
-        opts,args = self.parse_options(arg_s,'yr:')
+        opts,args = self.parse_options(arg_s,'ys:r:')
 
         if not args:
             raise UsageError('Missing filename, URL, input history range, '
                              'or macro.')
 
         contents = self.shell.find_user_code(args)
 
+        if 's' in opts:
+            blocks, not_found = extract_symbols(contents, opts['s'])
+            if len(not_found) == 1:
+                warn('The symbol `%s` was not found' % not_found[0])
+            elif len(not_found) > 1:
+                warn('The symbols %s were not found' % get_text_list(not_found,
+                                                                     wrap_item_with='`')
+                )
+
+            contents = '\n'.join(blocks)
+
         if 'r' in opts:
             ranges = opts['r'].replace(',', ' ')
             lines = contents.split('\n')
@@ -248,7 +316,6 @@ def load(self, arg_s):
 
         l = len(contents)
 
-
         # 200 000 is ~ 2500 full 80 caracter lines
         # so in average, more than 5000 lines
         if l > 200000 and 'y' not in opts:

IPython/core/tests/test_magic.py

@@ -60,6 +60,28 @@ def test_extract_code_ranges():
     actual = list(code.extract_code_ranges(instr))
     nt.assert_equal(actual, expected)
 
+
+def test_extract_symbols():
+    source = """import foo\na = 10\ndef b():\n    return 42\n\n\nclass A: pass\n\n\n"""
+    symbols_args = ["a", "b", "A", "A,b", "A,a", "z"]
+    expected = [([], ['a']),
+                (["def b():\n    return 42\n"], []),
+                (["class A: pass\n"], []),
+                (["class A: pass\n", "def b():\n    return 42\n"], []),
+                (["class A: pass\n"], ['a']),
+                ([], ['z'])]
+    for symbols, exp in zip(symbols_args, expected):
+        nt.assert_equal(code.extract_symbols(source, symbols), exp)
+
+
+def test_extract_symbols_ignores_non_python_code():
+    source = ("=begin A Ruby program :)=end\n"
+              "def hello\n"
+              "puts 'Hello world'\n"
+              "end")
+    nt.assert_equal(code.extract_symbols(source, "hello"), [])
+
+
 def test_rehashx():
     # clear up everything
     _ip = get_ipython()

IPython/utils/text.py

@@ -724,3 +724,32 @@ def columnize(items, separator='  ', displaywidth=80):
     fmatrix = [filter(None, x) for x in matrix]
     sjoin = lambda x : separator.join([ y.ljust(w, ' ') for y, w in zip(x, info['columns_width'])])
     return '\n'.join(map(sjoin, fmatrix))+'\n'
+
+
+def get_text_list(list_, last_word='and', wrap_item_with=""):
+    """
+    Return a string with a natural enumeration of items
+
+    >>> get_text_list(['a', 'b', 'c', 'd'])
+    'a, b, c and d'
+    >>> get_text_list(['a', 'b', 'c'], 'or')
+    'a, b or c'
+    >>> get_text_list(['a', 'b'], 'or')
+    'a or b'
+    >>> get_text_list(['a'])
+    'a'
+    >>> get_text_list([])
+    ''
+    >>> get_text_list(['a', 'b'], wrap_item_with="`")
+    '`a` and `b`'
+    """
+    if len(list_) == 0:
+        return ''
+    if wrap_item_with:
+        list_ = ['%s%s%s' % (wrap_item_with, item, wrap_item_with) for
+                 item in list_]
+    if len(list_) == 1:
+        return list_[0]
+    return '%s %s %s' % (
+        ', '.join(i for i in list_[:-1]),
+        last_word, list_[-1])