diff --git a/modules/base/pymod/table.py b/modules/base/pymod/table.py index 4bd1a9999738a8814554997a7e7ceb6d4e7c6cfe..097570b93e9b392a19d0b225a8e0fc322049db7e 100644 --- a/modules/base/pymod/table.py +++ b/modules/base/pymod/table.py @@ -601,13 +601,38 @@ class Table(object): return cPickle.load(stream) @staticmethod - def Load(stream_or_filename, format='ost', sep=','): + def _GuessFormat(filename): + try: + filename = filename.name + except AttributeError, e: + pass + if filename.endswith('.csv'): + return 'csv' + elif filename.endswith('.pickle'): + return 'pickle' + else: + return 'ost' + + + @staticmethod + def Load(stream_or_filename, format='auto', sep=','): """ Load table from stream or file with given name. - By default, the file format is *ost* (see below) and is *not* automatically - determined (e.g. from file extension). Thus, it *format* must be specified - for reading other file formats. + By default, the file format is set to *auto*, which tries to guess the file + format from the file extension. The following file extensions are + recognized: + + ============ ====================== + extension recognized format + ============ ====================== + .csv comma separated values + .pickle pickled byte stream + <all others> ost-specific format + ============ ====================== + + Thus, *format* must be specified for reading file with different filename + extensions. The following file formats are understood: @@ -643,6 +668,9 @@ class Table(object): :returns: A new :class:`Table` instance """ format=format.lower() + if format=='auto': + format = Table._GuessFormat(stream_or_filename) + if format=='ost': return Table._LoadOST(stream_or_filename) if format=='csv': diff --git a/modules/base/tests/test_table.py b/modules/base/tests/test_table.py index 4f8fdd505ca1c9765a0fc4fef40e7d2809ef007a..93ca61dd61751711da2a95a61e4d9f7b876b9b78 100644 --- a/modules/base/tests/test_table.py +++ b/modules/base/tests/test_table.py @@ -720,6 +720,53 @@ class TestTable(unittest.TestCase): tab.Sort('third', '+') self.CompareDataFromDict(tab, {'first': [None,'foo','x'], 'second': [9,None,3], 'third': [3.3,2.2,None]}) + def testGuessFormat(self): + self.assertEqual(Table._GuessFormat('table_test.csv'), 'csv') + self.assertEqual(Table._GuessFormat('table_test.pickle'), 'pickle') + self.assertEqual(Table._GuessFormat('table_test.tab'), 'ost') + self.assertEqual(Table._GuessFormat('table_test.ost'), 'ost') + self.assertEqual(Table._GuessFormat('table_test.xyz'), 'ost') + + def testSaveLoadTableAutoFormat(self): + tab = self.CreateTestTable() + self.CompareDataFromDict(tab, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + + # write to disc + tab.Save("saveloadtable_filename_out.csv", format='csv') + tab.Save("saveloadtable_filename_out.tab", format='ost') + tab.Save("saveloadtable_filename_out.pickle", format='pickle') + + # read from disc: csv + in_stream_csv = open("saveloadtable_stream_out.csv", 'r') + tab_loaded_stream_csv = Table.Load(in_stream_csv) + in_stream_csv.close() + tab_loaded_fname_csv = Table.Load('saveloadtable_filename_out.csv') + + # check content: csv + self.CompareDataFromDict(tab_loaded_stream_csv, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + self.CompareDataFromDict(tab_loaded_fname_csv, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + + # read from disc: pickle + in_stream_pickle = open("saveloadtable_stream_out.pickle", 'rb') + tab_loaded_stream_pickle = Table.Load(in_stream_pickle) + in_stream_pickle.close() + tab_loaded_fname_pickle = Table.Load('saveloadtable_filename_out.pickle') + + # check content: pickle + self.CompareDataFromDict(tab_loaded_stream_pickle, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + self.CompareDataFromDict(tab_loaded_fname_pickle, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + + # read from disc: ost + in_stream_ost = open("saveloadtable_stream_out.tab", 'rb') + tab_loaded_stream_ost = Table.Load(in_stream_ost) + in_stream_ost.close() + tab_loaded_fname_ost = Table.Load('saveloadtable_filename_out.tab') + + # check content: ost + self.CompareDataFromDict(tab_loaded_stream_ost, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + self.CompareDataFromDict(tab_loaded_fname_ost, {'first': ['x','foo',None], 'second': [3,None,9], 'third': [None,2.2,3.3]}) + + def testLoadTableOSTUnknownType(self): self.assertRaises(ValueError, Table.Load, os.path.join('testfiles','ost-table-unknown-type.tab'))