Show More
@@ -437,6 +437,21 b' class TestAstTransform(unittest.TestCase):' | |||
|
437 | 437 | ip.user_ns['n'] = 55 |
|
438 | 438 | with tt.AssertNotPrints('-55'): |
|
439 | 439 | ip.run_cell('print (n)') |
|
440 | ||
|
441 | def test_timeit(self): | |
|
442 | called = set() | |
|
443 | def f(x): | |
|
444 | called.add(x) | |
|
445 | ip.push({'f':f}) | |
|
446 | ||
|
447 | with tt.AssertPrints("best of "): | |
|
448 | ip.run_line_magic("timeit", "-n1 f(1)") | |
|
449 | self.assertEqual(called, set([-1])) | |
|
450 | called.clear() | |
|
451 | ||
|
452 | with tt.AssertPrints("best of "): | |
|
453 | ip.run_cell_magic("timeit", "-n1 f(2)", "f(3)") | |
|
454 | self.assertEqual(called, set([-2, -3])) | |
|
440 | 455 | |
|
441 | 456 | class IntegerWrapper(ast.NodeTransformer): |
|
442 | 457 | """Wraps all integers in a call to Integer()""" |
@@ -453,6 +468,7 b' class TestAstTransform2(unittest.TestCase):' | |||
|
453 | 468 | self.calls = [] |
|
454 | 469 | def Integer(*args): |
|
455 | 470 | self.calls.append(args) |
|
471 | return args | |
|
456 | 472 | ip.push({"Integer": Integer}) |
|
457 | 473 | |
|
458 | 474 | def tearDown(self): |
@@ -462,6 +478,21 b' class TestAstTransform2(unittest.TestCase):' | |||
|
462 | 478 | def test_run_cell(self): |
|
463 | 479 | ip.run_cell("n = 2") |
|
464 | 480 | self.assertEqual(self.calls, [(2,)]) |
|
481 | ||
|
482 | def test_timeit(self): | |
|
483 | called = set() | |
|
484 | def f(x): | |
|
485 | called.add(x) | |
|
486 | ip.push({'f':f}) | |
|
487 | ||
|
488 | with tt.AssertPrints("best of "): | |
|
489 | ip.run_line_magic("timeit", "-n1 f(1)") | |
|
490 | self.assertEqual(called, set([(1,)])) | |
|
491 | called.clear() | |
|
492 | ||
|
493 | with tt.AssertPrints("best of "): | |
|
494 | ip.run_cell_magic("timeit", "-n1 f(2)", "f(3)") | |
|
495 | self.assertEqual(called, set([(2,), (3,)])) | |
|
465 | 496 | |
|
466 | 497 | class ErrorTransformer(ast.NodeTransformer): |
|
467 | 498 | """Throws an error when it sees a number.""" |
General Comments 0
You need to be logged in to leave comments.
Login now