import nose
import types
import os
import pandas as pd

from collections import defaultdict
from nose.tools import make_decorator
from IPython.core.magic import Magics, magics_class, cell_magic


class score(object):
    """Decorator for marking the problem an autograder test corresponds
    to, as well as the number of points that should be earned if the
    test passes.

    The `grades` and `max_grades` dictionaries should be reset (using
    `score.reset()`) before the autograding code is actually run. This
    will happen automatically if the `%%autograde` magic is used.

    """

    grades = defaultdict(float)
    max_grades = defaultdict(float)

    def __init__(self, problem, points):
        self.problem = problem
        self.points = points
        self.max_grades[self.problem] += self.points

    def __call__(self, f):
        def wrapped_f(*args):
            f(*args)
            self.grades[self.problem] += self.points
        return make_decorator(f)(wrapped_f)

    @classmethod
    def reset(cls):
        cls.grades = defaultdict(float)
        cls.max_grades = defaultdict(float)

    @classmethod
    def as_dataframe(cls):
        df = pd.DataFrame({
            'score': cls.grades,
            'max_score': cls.max_grades
        }, columns=['score', 'max_score'])
        df.fillna(0, inplace=True)
        return df

    @classmethod
    def as_dict(cls):
        df = cls.as_dataframe()
        return df.T.to_dict()


@magics_class
class NoseGraderMagic(Magics):

    @cell_magic
    def autograde(self, line, cell):
        # get the ipython user session, import the score decorator,
        # and then run the autograding code (which will load the tests
        # into the namespace)
        ip = get_ipython()
        ip.run_cell("from autograder import score; score.reset()")
        ip.run_cell(cell)

        # create the test module for nose
        test_module = types.ModuleType('test_module')
        test_module.__dict__.update(ip.user_ns)

        # change the test name template to use "grade" instead of
        # "test" (which is the default)
        env = os.environ
        env['NOSE_TESTMATCH'] = r'(?:^|[\b_\.%s-])[Gg]rade' % os.sep

        # load user config files
        cfg_files = nose.config.all_config_files()

        # create a plugin manager
        mgr = nose.plugins.manager.DefaultPluginManager()

        # create the nose configuration object, and load the tests
        config = nose.config.Config(env=env, files=cfg_files, plugins=mgr)
        loader = nose.loader.TestLoader(config=config)
        tests = loader.loadTestsFromModule(test_module)

        # run the tests
        nose.core.TestProgram(
            argv=["autograder", "--verbose"],
            suite=tests, exit=False, config=config)

        # create a pandas dataframe of the scores, and return it
        return test_module.__dict__['score'].as_dataframe()


def load_ipython_extension(ipython):
    ipython.register_magics(NoseGraderMagic)

@@ -132,56 +132,159 @@
 rom 
-IPython.core.magic import Magics, magics_class, 
+nose.plugins.attrib import attr, AttributeSelector%0Afrom IPython.core.magic import Magics, magics_class%0Afrom IPython.core.magic import line_magic, line_
 cell
@@ -653,20 +653,16 @@
  the %60%25%25
-auto
 grade%60 m
@@ -1055,16 +1055,43 @@
  return 
+attr(problem=self.problem)(
 make_dec
@@ -1110,16 +1110,17 @@
 apped_f)
+)
 %0A%0A    @c
@@ -1631,16 +1631,21 @@
 :%0A%0A    @
+line_
 cell_mag
@@ -1655,20 +1655,16 @@
     def 
-auto
 grade(se
@@ -1677,16 +1677,21 @@
 ne, cell
+=None
 ):%0A     
@@ -1697,196 +1697,194 @@
     
-# get the ipython user session, import the score decorator,%0A        # and then run the autograding code (which will load the tests%0A        # into the namespace)%0A        ip = get_ipython()%0A
+ip = get_ipython()%0A        if cell is not None:%0A            ip.run_cell(cell)%0A        else:%0A            with open(self.autograder_path, %22r%22) as fh:%0A                code = fh.read()%0A%0A    
     
@@ -1932,26 +1932,15 @@
 core
-; score.reset()%22)%0A
+%22)%0A    
     
@@ -1956,19 +1956,19 @@
 n_cell(c
-ell
+ode
 )%0A%0A     
@@ -2427,16 +2427,124 @@
 manager%0A
+        plug = AttributeSelector()%0A        plug.attribs = %5B%5B(%22problem%22, line)%5D%5D%0A        plug.enabled = True%0A
         
@@ -2591,16 +2591,43 @@
 Manager(
+%0A            plugins=%5Bplug%5D
 )%0A%0A     
@@ -3106,59 +3106,227 @@
     
-return test_module.__dict__%5B'score'%5D.as_dataframe()
+if 'score' in test_module.__dict__:%0A            df = test_module.__dict__%5B'score'%5D.as_dataframe()%0A            return df.ix%5B%5Bline%5D%5D%0A%0A    @line_magic%0A    def set_autograder(self, line):%0A        self.autograder_path = line
 %0A%0A%0Ad

Update autograder magic with line magic support and test filtering