def test_no_extra_needed(self): command_list = ['df.reset_index()', 'df.head()'] self.assertEqual( set(get_modules_and_shortcuts(command_list)), { ('pandas', 'pd'), ('dateutil', 'dateutil'), } )
def test_get_extra_import_all_needed(self): command_list = [ 'pl.plot(df.x)', 'sns.distplot(df.x)', 'scp.stats.norm(1, 1)', 'np.random.randn(1)' ] self.assertEqual( set(get_modules_and_shortcuts(command_list)), { ('dateutil', 'dateutil'), ('pandas', 'pd'), ('scipy', 'scp'), ('pylab', 'pl'), ('seaborn', 'sns'), ('numpy', 'np'), }, )